Support for passing in known classes for given oids.

Can be used if there is no 1-to-1 mapping from old to new classes.

CHANGES.txt

@@ -4,7 +4,9 @@ Changes
 0.6 (unreleased)
 ----------------
 
-- ...
+- Support for passing in known classes for given oids.
+  Can be used if there is no 1-to-1 mapping from old to new classes.
+
 
 0.5 (2010-10-07)
 ----------------

src/zodbupdate/main.py

@@ -15,6 +15,8 @@
 import ZODB.FileStorage
 import ZODB.config
 import ZODB.serialize
+import ZODB.utils
+import json
 import logging
 import optparse
 import pkg_resources
@@ -57,7 +59,9 @@
 parser.add_option(
     "--pack", action="store_true", dest="pack",
     help="pack the storage when done. use in conjunction of -c if you have blobs storage")
-
+parser.add_option(
+    "-k", "--known-classes",
+    help="file containing a mapping from oids to module/class tuples in JSON format")
 
 
 class DuplicateFilter(object):
@@ -108,6 +112,13 @@ def main():
     if options.oid:
         start_at = options.oid
 
+    class_by_oid = {}
+    if options.known_classes:
+        classes_file = open(options.known_classes, 'r')
+        class_by_oid = {ZODB.utils.repr_to_oid(key): value
+                        for key, value in json.load(classes_file).items()}
+        classes_file.close()
+
     rename_rules = {}
     for entry_point in pkg_resources.iter_entry_points('zodbupdate'):
         rules = entry_point.load()
@@ -119,6 +130,7 @@ def main():
         storage,
         dry=options.dry_run,
         renames=rename_rules,
+        class_by_oid=class_by_oid,
         start_at=start_at,
         debug=options.debug,
         pickler_name=options.pickler)

src/zodbupdate/serialize.py

@@ -144,21 +144,25 @@ class ObjectRenamer(object):
     - in class information (first pickle of the record).
     """
 
-    def __init__(self, changes, pickler_name='C'):
+    def __init__(self, changes, class_by_oid={}, pickler_name='C'):
         self.__added = dict()
         self.__changes = dict()
         for old, new in changes.iteritems():
             self.__changes[tuple(old.split(' '))] = tuple(new.split(' '))
+        self.__class_by_oid = class_by_oid
         self.__changed = False
         self.__pickler_name = pickler_name
 
-    def __update_symb(self, symb_info):
+    def __update_symb(self, symb_info, oid=None):
         """This method look in a klass or symbol have been renamed or
         not. If the symbol have not been renamed explicitly, it's
         loaded and its location is checked to see if it have moved as
         well.
         """
-        if symb_info in self.__changes:
+        if oid in self.__class_by_oid:
+            self.__changed = True
+            return self.__class_by_oid[oid]
+        elif symb_info in self.__changes:
             self.__changed = True
             return self.__changes[symb_info]
         else:
@@ -195,15 +199,21 @@ def __persistent_load(self, reference):
         """
         if isinstance(reference, tuple):
             oid, klass_info = reference
-            if isinstance(klass_info, tuple):
-                klass_info = self.__update_symb(klass_info)
+            if oid in self.__class_by_oid:
+                self.__changed = True
+                klass_info = find_global(*self.__class_by_oid[oid], Broken=ZODBBroken)
+            elif isinstance(klass_info, tuple):
+                klass_info = self.__update_symb(klass_info, oid=oid)
             return ZODBReference((oid, klass_info))
         if isinstance(reference, list):
             mode, information = reference
             if mode == 'm':
                 database_name, oid, klass_info = information
-                if isinstance(klass_info, tuple):
-                    klass_info = self.__update_symb(klass_info)
+                if oid in self.__class_by_oid:
+                    self.__changed = True
+                    klass_info = find_global(*self.__class_by_oid[oid], Broken=ZODBBroken)
+                elif isinstance(klass_info, tuple):
+                    klass_info = self.__update_symb(klass_info, oid=oid)
                 return ZODBReference(['m', (database_name, oid, klass_info)])
         return ZODBReference(reference)
 
@@ -230,7 +240,7 @@ def __pickler(self, output_file):
         pickler.persistent_id = self.__persistent_id
         return pickler
 
-    def __update_class_meta(self, class_meta):
+    def __update_class_meta(self, class_meta, oid):
         """Update class information, which can contain information
         about a renamed class.
         """
@@ -242,10 +252,10 @@ def __update_class_meta(self, class_meta):
                     u'Warning: Missing factory for %s' % u' '.join(symb_info))
                 return (symb_info, args)
             elif isinstance(symb, tuple):
-                return self.__update_symb(symb), args
+                return self.__update_symb(symb, oid=oid), args
         return class_meta
 
-    def rename(self, input_file):
+    def rename(self, input_file, oid):
         """Take a ZODB record (as a file object) as input. We load it,
         replace any reference to renamed class we know of. If any
         modification are done, we save the record again and return it,
@@ -257,7 +267,7 @@ def rename(self, input_file):
         class_meta = unpickler.load()
         data = unpickler.load()
 
-        class_meta = self.__update_class_meta(class_meta)
+        class_meta = self.__update_class_meta(class_meta, oid)
 
         if not (self.__changed or
                 (hasattr(unpickler, 'need_repickle') and

src/zodbupdate/update.py

@@ -30,12 +30,12 @@
 class Updater(object):
     """Update class references for all current objects in a storage."""
 
-    def __init__(self, storage, dry=False, renames=None,
+    def __init__(self, storage, dry=False, renames=None, class_by_oid={},
                  start_at='0x00', debug=False, pickler_name='C'):
         self.dry = dry
         self.storage = storage
         self.processor = zodbupdate.serialize.ObjectRenamer(
-            renames or {}, pickler_name)
+            renames or {}, class_by_oid=class_by_oid, pickler_name=pickler_name)
         self.start_at = start_at
         self.debug = debug
 
@@ -62,7 +62,7 @@ def __call__(self):
             for oid, serial, current in self.records:
                 logger.debug('Processing OID %s' % ZODB.utils.oid_repr(oid))
 
-                new = self.processor.rename(current)
+                new = self.processor.rename(current, oid)
                 if new is None:
                     continue