Support for extension codes. (By accident I checked in the tests first.)

1 files changed, 59 insertions, 0 deletions

diff --git a/Lib/copy_reg.py b/Lib/copy_reg.py
index 8a3550a..1276564 100644
--- a/Lib/copy_reg.py
+++ b/Lib/copy_reg.py
@@ -70,3 +70,62 @@ def _reduce(self):
         return _reconstructor, args, dict
     else:
         return _reconstructor, args
+
+# A registry of extension codes.  This is an ad-hoc compression
+# mechanism.  Whenever a global reference to <module>, <name> is about
+# to be pickled, the (<module>, <name>) tuple is looked up here to see
+# if it is a registered extension code for it.  Extension codes are
+# universal, so that the meaning of a pickle does not depend on
+# context.  (There are also some codes reserved for local use that
+# don't have this restriction.)  Codes are positive ints; 0 is
+# reserved.
+
+extension_registry = {}                 # key -> code
+inverted_registry = {}                  # code -> key
+extension_cache = {}                    # code -> object
+
+def add_extension(module, name, code):
+    """Register an extension code."""
+    code = int(code)
+    if not 1 <= code < 0x7fffffff:
+        raise ValueError, "code out of range"
+    key = (module, name)
+    if (extension_registry.get(key) == code and
+        inverted_registry.get(code) == key):
+        return # Redundant registrations are benign
+    if key in extension_registry:
+        raise ValueError("key %s is already registered with code %s" %
+                         (key, extension_registry[key]))
+    if code in inverted_registry:
+        raise ValueError("code %s is already in use for key %s" %
+                         (code, inverted_registry[code]))
+    extension_registry[key] = code
+    inverted_registry[code] = key
+
+def remove_extension(module, name, code):
+    """Unregister an extension code.  For testing only."""
+    key = (module, name)
+    if (extension_registry.get(key) != code or
+        inverted_registry.get(code) != key):
+        raise ValueError("key %s is not registered with code %s" %
+                         (key, code))
+    del extension_registry[key]
+    del inverted_registry[code]
+    if code in extension_cache:
+        del extension_cache[code]
+
+def clear_extension_cache():
+    extension_cache.clear()
+
+# Standard extension code assignments
+
+# Reserved ranges
+
+# First  Last Count  Purpose
+#     1   127   127  Reserved for Python standard library
+#   128   191    64  Reserved for Zope 3
+#   192   239    48  Reserved for 3rd parties
+#   240   255    16  Reserved for private use (will never be assigned)
+#   256   Inf   Inf  Reserved for future assignment
+
+# Extension codes are assigned by the Python Software Foundation.