Address review comments: Only update one cell, then break out of the …

…loop; minor change to comments; fix a test.
python · ericvsmith · Sep 25, 2024 · Sep 24, 2024 · Sep 24, 2024 · Sep 24, 2024
commit 1c60e6b77e90f894324f27641f1b83f68eca5356
@@ -1222,21 +1222,24 @@ def _get_slots(cls):


 def _update_func_cell_for__class__(f, oldcls, newcls):
+    # Returns True if we update a cell, else False.
    if f is None:
        # f will be None in the case of a property where not all of
        # fget, fset, and fdel are used.  Nothing to do in that case.
-        return
+        return False
    try:
        idx = f.__code__.co_freevars.index("__class__")
    except ValueError:
        # This function doesn't reference __class__, so nothing to do.
-        return
+        return False
    # Fix the cell to point to the new class, if it's already pointing
    # at the old class.  I'm not convinced that the "is oldcls" test
    # is needed, but other than performance can't hurt.
    closure = f.__closure__[idx]
    if closure.cell_contents is oldcls:
        closure.cell_contents = newcls
+        return True
+    return False


 def _add_slots(cls, is_frozen, weakref_slot):
@@ -1295,17 +1298,21 @@ def _add_slots(cls, is_frozen, weakref_slot):
    # Fix up any closures which reference __class__.  This is used to
    # fix zero argument super so that it points to the correct class
    # (the newly created one, which we're returning) and not the
-    # original class.
+    # original class.  We can break out of this loop as soon as we
+    # make an update, since all closures for a class will share a
+    # given cell.
    for member in newcls.__dict__.values():
        # If this is a wrapped function, unwrap it.
        member = inspect.unwrap(member)

        if isinstance(member, types.FunctionType):
-            _update_func_cell_for__class__(member, cls, newcls)
+            if _update_func_cell_for__class__(member, cls, newcls):
+                break
        elif isinstance(member, property):
-            _update_func_cell_for__class__(member.fget, cls, newcls)
-            _update_func_cell_for__class__(member.fset, cls, newcls)
-            _update_func_cell_for__class__(member.fdel, cls, newcls)
+            if (_update_func_cell_for__class__(member.fget, cls, newcls)
+                or _update_func_cell_for__class__(member.fset, cls, newcls)
+                or _update_func_cell_for__class__(member.fdel, cls, newcls)):
+                break

    return newcls


@@ -4884,6 +4884,7 @@ class A:
            def _get_foo(slf):
                self.assertIs(__class__, type(slf))
                self.assertIs(__class__, slf.__class__)
+                return __class__

            def _set_foo(slf, value):
                self.assertIs(__class__, type(slf))
@@ -4981,11 +4982,11 @@ def cls(self):
        B = dataclass(slots=True)(A)
        self.assertIs(B().cls(), B)

-        # This is probably undesirable behavior, but is a function of
-        # how modifying __class__ in the closure works.  I'm not sure
-        # this should be tested or not: I don't really want to
-        # guarantee this behavior, but I don't want to lose the point
-        # that this is how it works.
+        # This is undesirable behavior, but is a function of how
+        # modifying __class__ in the closure works.  I'm not sure this
+        # should be tested or not: I don't really want to guarantee
+        # this behavior, but I don't want to lose the point that this
+        # is how it works.

        # The underlying class is "broken" by changing its __class__
        # in A.foo() to B.  This normally isn't a problem, because no