Ticket #198: class_definition_reverse_lookup.diff

pyamf/tests/test_amf3.py

@@ -728 +728 @@
         self.assertFalse(is_ref)
         self.assertTrue(isinstance(class_def, amf3.ClassDefinition))
         self.assertTrue(class_def.getClass(), Spam)
+        self.assertTrue(class_def.klass, Spam)
         self.assertTrue(class_def.name, 'abc.xyz')
         self.assertEquals(class_def.encoding, amf3.ObjectEncoding.STATIC)
 
         self.assertTrue(class_def in self.context.classes)
+        self.assertTrue(id(class_def) in self.context.rev_classes)
 
         self.context.classes.remove(class_def)
         self.buf.write('\x0fabc.xyz')

pyamf/amf3.py

@@ -658 +658 @@
         pyamf.BaseContext.clear(self)
 
         self.strings = []
+
         self.classes = []
+        self.rev_classes = {}
         self.class_defs = {}
+        self.rev_class_defs = {}
+
         self.legacy_xml = []
+        self.rev_legacy_xml = {}
 
     def getString(self, ref):
         """
@@ -676 +681 @@
         try:
             return self.strings[ref]
         except IndexError:
-            raise pyamf.ReferenceError("String reference %d not found" % ref)
+            raise pyamf.ReferenceError, "String reference %d not found" % ref
 
     def getStringReference(self, s):
         """
@@ -691 +696 @@
         try:
             return self.strings.index(s)
         except ValueError:
-            raise pyamf.ReferenceError("Reference for string %r not found" % s)
+            raise pyamf.ReferenceError, "Reference for string %r not found" % s
 
     def addString(self, s):
         """
@@ -705 +710 @@
         @rtype: C{int}
         @return: The reference index.
         """
+        if not isinstance(s, basestring):
+            raise TypeError
+
         if len(s) == 0:
             # do not store empty string references
             raise ValueError, "Cannot store a reference to an empty string"
@@ -728 +736 @@
         try:
             return self.classes[ref]
         except IndexError:
-            raise pyamf.ReferenceError("Class reference %d not found" % ref)
+            raise pyamf.ReferenceError, "Class reference %d not found" % ref
 
     def getClassDefinitionReference(self, class_def):
         """
         Return class definition reference.
 
-        @type class_def: L{ClassDefinition}
+        @type class_def: L{ClassDefinition} or C{instance} or C{class}
         @param class_def: The class definition reference to be found.
         @raise ReferenceError: The reference could not be found.
         @return: The reference to C{class_def}.
         @rtype: C{int}
         """
-        if not isinstance(class_def, ClassDefinition) and isinstance(class_def, (type, types.ClassType)):
+        if not isinstance(class_def, ClassDefinition):
+            if isinstance(class_def, (type, types.ClassType)):
+                try:
+                    return self.rev_class_defs[class_def]
+                except KeyError:
+                    raise pyamf.ReferenceError("Reference for class definition for %s not found" %
+                        class_def)
+            elif isinstance(class_def, (types.InstanceType, types.ObjectType)):
+                try:
+                    return self.class_defs[class_def.__class__]
+                except KeyError:
+                    raise pyamf.ReferenceError("Reference for class definition for %s not found" %
+                        class_def.__class__)
+
+            raise TypeError, 'unable to determine class for %r' % class_def
+        else:
             try:
-                return self.class_defs[class_def.__class__]
-            except KeyError:
-                raise pyamf.ReferenceError("Reference for class %s not found" %
-                    class_def.__class__)
+                return self.rev_class_defs[id(class_def)]
+            except ValueError:
+                raise pyamf.ReferenceError, "Reference for class %s not found" % kl
 
-        try:
-            return self.classes.index(class_def)
-        except ValueError:
-            raise pyamf.ReferenceError("Reference for class %s not found" %
-                class_def.__class__)
-
     def addClassDefinition(self, class_def):
         """
         Creates a reference to C{class_def}.
         """
         try:
-            return self.classes.index(class_def)
-        except ValueError:
+            return self.rev_class_defs[id(class_def)]
+        except KeyError:
             self.classes.append(class_def)
+            idx = len(self.classes) - 1
+
+            self.rev_classes[id(class_def)] = idx
             self.class_defs[class_def.__class__] = class_def
+            self.rev_class_defs[id(class_def.__class__)] = idx
 
-            return len(self.classes) - 1
+            return idx
 
     def getLegacyXML(self, ref):
         """
@@ -793 +813 @@
         @rtype: C{int}
         """
         try:
-            return self.legacy_xml.index(doc)
-        except ValueError:
+            return self.rev_legacy_xml[id(doc)]
+        except KeyError:
             raise pyamf.ReferenceError, "Reference for document %r not found" % doc
 
     def addLegacyXML(self, doc):
@@ -810 +830 @@
         @return: The reference to C{doc}.
         """
         try:
-            return self.legacy_xml.index(doc)
-        except ValueError:
+            return self.rev_legacy_xml[id(doc)]
+        except KeyError:
             self.legacy_xml.append(doc)
 
-            return len(self.legacy_xml) - 1
+            idx = len(self.legacy_xml) - 1
+            self.rev_legacy_xml[id(doc)] = idx
 
+            return idx
+
     def __copy__(self):
         return self.__class__()
 
@@ -1161 +1184 @@
         # Unsupported types go first
         ((types.BuiltinFunctionType, types.BuiltinMethodType,),
             "writeUnsupported"),
-        ((lambda x: x == pyamf.Undefined,), "writeUndefined"),
+        ((lambda x: x is pyamf.Undefined,), "writeUndefined"),
         ((bool,), "writeBoolean"),
         ((int,long), "writeInteger"),
         ((float,), "writeNumber"),
@@ -1562 +1585 @@
 
             self._writeInteger(ref | class_def.encoding << 2 | REFERENCE_BIT << 1 | REFERENCE_BIT)
             self._writeString(class_def.name)
+        else:
+            class_def = self._getClassDefinition(obj)
 
         if class_def.encoding in (ObjectEncoding.EXTERNAL, ObjectEncoding.PROXY):
             obj.__writeamf__(DataOutput(self))