Python: CG trace: Apply better_compare_for_dataclass to all

RasmusWL · RasmusWL · commit b86ca19264bb · 2020-07-21T23:37:33.000+02:00
diff --git a/python/tools/recorded-call-graph-metrics/src/cg_trace/bytecode_reconstructor.py b/python/tools/recorded-call-graph-metrics/src/cg_trace/bytecode_reconstructor.py
@@ -5,6 +5,8 @@
 from types import FrameType
 from typing import Any, List
 
+from cg_trace.utils import better_compare_for_dataclass
+
 LOGGER = logging.getLogger(__name__)
 
 # See https://docs.python.org/3/library/dis.html#python-bytecode-instructions for
@@ -18,6 +20,7 @@ class BytecodeExpr:
     """
 
 
+@better_compare_for_dataclass
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeConst(BytecodeExpr):
     """FOR LOAD_CONST"""
@@ -28,6 +31,7 @@ def __str__(self):
         return repr(self.value)
 
 
+@better_compare_for_dataclass
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeVariableName(BytecodeExpr):
     name: str
@@ -36,6 +40,7 @@ def __str__(self):
         return self.name
 
 
+@better_compare_for_dataclass
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeAttribute(BytecodeExpr):
     attr_name: str
@@ -45,6 +50,7 @@ def __str__(self):
         return f"{self.object}.{self.attr_name}"
 
 
+@better_compare_for_dataclass
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeSubscript(BytecodeExpr):
     key: BytecodeExpr
@@ -54,6 +60,7 @@ def __str__(self):
         return f"{self.object}[{self.key}]"
 
 
+@better_compare_for_dataclass
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeTuple(BytecodeExpr):
     elements: List[BytecodeExpr]
@@ -67,6 +74,7 @@ def __str__(self):
         return f"({elements_formatted})"
 
 
+@better_compare_for_dataclass
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeList(BytecodeExpr):
     elements: List[BytecodeExpr]
@@ -80,6 +88,7 @@ def __str__(self):
         return f"[{elements_formatted}]"
 
 
+@better_compare_for_dataclass
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeCall(BytecodeExpr):
     function: BytecodeExpr
@@ -88,6 +97,7 @@ def __str__(self):
         return f"{self.function}()"
 
 
+@better_compare_for_dataclass
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeUnknown(BytecodeExpr):
     opname: str
@@ -96,6 +106,7 @@ def __str__(self):
         return f"<{self.opname}>"
 
 
+@better_compare_for_dataclass
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeMakeFunction(BytecodeExpr):
     """For MAKE_FUNCTION opcode"""
@@ -106,6 +117,7 @@ def __str__(self):
         return f"<MAKE_FUNCTION>(qualified_name={self.qualified_name})>"
 
 
+@better_compare_for_dataclass
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class SomethingInvolvingScaryBytecodeJump(BytecodeExpr):
     opname: str
diff --git a/python/tools/recorded-call-graph-metrics/src/cg_trace/tracer.py b/python/tools/recorded-call-graph-metrics/src/cg_trace/tracer.py
@@ -6,6 +6,7 @@
 from typing import Any, Optional, Tuple
 
 from cg_trace.bytecode_reconstructor import BytecodeExpr, expr_from_frame
+from cg_trace.utils import better_compare_for_dataclass
 
 LOGGER = logging.getLogger(__name__)
 
@@ -75,28 +76,6 @@ def hash_key(frame: FrameType) -> Tuple[str, int, int]:
         )
 
 
-def better_compare_for_dataclass(cls):
-    """When dataclass is used with `order=True`, the comparison methods is only implemented for
-    objects of the same class. This decorator extends the functionality to compare class
-    name if used against other objects.
-    """
-    for op in [
-        "__lt__",
-        "__le__",
-        "__gt__",
-        "__ge__",
-    ]:
-        old = getattr(cls, op)
-
-        def new(self, other):
-            if type(self) == type(other):
-                return old(self, other)
-            return getattr(str, op)(self.__class__.__name__, other.__class__.__name__)
-
-        setattr(cls, op, new)
-    return cls
-
-
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class Callee:
     pass
diff --git a/python/tools/recorded-call-graph-metrics/src/cg_trace/utils.py b/python/tools/recorded-call-graph-metrics/src/cg_trace/utils.py
@@ -0,0 +1,20 @@
+def better_compare_for_dataclass(cls):
+    """When dataclass is used with `order=True`, the comparison methods is only implemented for
+    objects of the same class. This decorator extends the functionality to compare class
+    name if used against other objects.
+    """
+    for op in [
+        "__lt__",
+        "__le__",
+        "__gt__",
+        "__ge__",
+    ]:
+        old = getattr(cls, op)
+
+        def new(self, other):
+            if type(self) == type(other):
+                return old(self, other)
+            return getattr(str, op)(self.__class__.__name__, other.__class__.__name__)
+
+        setattr(cls, op, new)
+    return cls
