--- a/contrib/memory.py	Tue Apr 11 14:54:12 2017 -0700
+++ b/contrib/memory.py	Tue Apr 11 14:54:12 2017 -0700
@@ -12,7 +12,6 @@
 '''
 
 from __future__ import absolute_import
-import atexit
 
 def memusage(ui):
     """Report memory usage of the current process."""
@@ -29,4 +28,4 @@
                             for k, v in result.iteritems()]) + "\n")
 
 def extsetup(ui):
-    atexit.register(memusage, ui)
+    ui.atexit(memusage, ui)