GH-100930: Implement dataclass code caching

Lib/dataclasses.py

@@ -8,7 +8,7 @@
 import itertools
 import abc
 import _thread
-from types import FunctionType, GenericAlias
+from types import CellType, FunctionType, GenericAlias
 
 
 __all__ = ['dataclass',
@@ -222,6 +222,9 @@ def __repr__(self):
 # https://bugs.python.org/issue33453 for details.
 _MODULE_IDENTIFIER_RE = re.compile(r'^(?:\s*(\w+)\s*\.)?\s*(\w+)')
 
+_PLACEHOLDER_PREFIX = "_field_"
+_PLACEHOLDER_RE = re.compile(rf"({_PLACEHOLDER_PREFIX}\w+)")
+
 # This function's logic is copied from "recursive_repr" function in
 # reprlib module to avoid dependency.
 def _recursive_repr(user_function):
@@ -281,6 +284,7 @@ class Field:
                  'metadata',
                  'kw_only',
                  '_field_type',  # Private: not to be used by user code.
+                 '_placeholder',
                  )
 
     def __init__(self, default, default_factory, init, repr, hash, compare,
@@ -298,6 +302,7 @@ def __init__(self, default, default_factory, init, repr, hash, compare,
                          types.MappingProxyType(metadata))
         self.kw_only = kw_only
         self._field_type = None
+        self._placeholder = None
 
     @_recursive_repr
     def __repr__(self):
@@ -410,36 +415,77 @@ def _fields_in_init_order(fields):
             tuple(f for f in fields if f.init and f.kw_only)
             )
 
-
 def _tuple_str(obj_name, fields):
     # Return a string representing each field of obj_name as a tuple
     # member.  So, if fields is ['x', 'y'] and obj_name is "self",
-    # return "(self.x,self.y)".
+    # return "(self._field_x,self._field_y)".
 
     # Special case for the 0-tuple.
     if not fields:
         return '()'
     # Note the trailing comma, needed if this turns out to be a 1-tuple.
-    return f'({",".join([f"{obj_name}.{f.name}" for f in fields])},)'
+    return f'({",".join([f"{obj_name}.{f._placeholder}" for f in fields])},)'
 
 
-def _create_fn(name, args, body, *, globals=None, locals=None,
-               return_type=MISSING):
+_code_cache = {}
+_code_cache_hits = 0
+
+def _create_fn(name, args, body, *, globals=None, locals=None):
     # Note that we may mutate locals. Callers beware!
     # The only callers are internal to this module, so no
     # worries about external callers.
     if locals is None:
         locals = {}
-    return_annotation = ''
-    if return_type is not MISSING:
-        locals['_return_type'] = return_type
-        return_annotation = '->_return_type'
     args = ','.join(args)
     body = '\n'.join(f'  {b}' for b in body)
 
     # Compute the text of the entire function.
-    txt = f' def {name}({args}){return_annotation}:\n{body}'
+    txt = f' def {name}({args}):\n{body}'
+
+    # exec is *really* slow, so we want to avoid using it whenever possible.
+    # Turns out, if we already have a code object that is the exact same as the
+    # one we want to create (except for the actual names of the fields), it's
+    # twice as fast to just copy the other code object, fix up the names, and
+    # create a function using the types.FunctionType constructor.
+
+    # In order for this to work, we need to be able to detect the names of the
+    # fields from the source. For this reason, it's *crucial* that anything
+    # building source strings uses field._placeholder, *not* field.name!
+
+    # This trick requires two passes:
+
+    # - Strip out the field names from the source text, and replace them with
+    #   something else (like numbers). This is used to search for (or create)
+    #   cached code with the same structure. This happens in _extract_fields.
 
+    # - Make a copy of the cached code object, and fix up all of the stripped
+    #   names in the consts, names, and varnames. This happens in _patch_fields.
+
+    # This is surprisingly effective, since we generate lots of functions that
+    # differ *only* in the names of their fields (we avoid something like 95% of
+    # the exec calls in test_dataclasses.py). Credit for this neat idea goes to
+    # https://github.com/dabeaz/dataklasses.
+
+    patches, key = _extract_fields(txt)
+    if key in _code_cache:
+        global _code_cache_hits
+        _code_cache_hits += 1
+        code = _code_cache[key]
+    else:
+        code = _code_cache[key] = _create_code(name, key, globals, locals)
+    # Build the closure:
+    closure = []
+    for freevar in code.co_freevars:
+        closure.append(CellType(locals[freevar]))
+    # Build the function:
+    return FunctionType(
+        code=_patch_fields(patches, code),
+        globals=globals or {},
+        closure=tuple(closure),
+    )
+
+
+def _create_code(name, txt, globals, locals):
     # Free variables in exec are resolved in the global namespace.
     # The global namespace we have is user-provided, so we can't modify it for
     # our purposes. So we put the things we need into locals and introduce a
@@ -448,7 +494,68 @@ def _create_fn(name, args, body, *, globals=None, locals=None,
     txt = f"def __create_fn__({local_vars}):\n{txt}\n return {name}"
     ns = {}
     exec(txt, globals, ns)
-    return ns['__create_fn__'](**locals)
+    return ns['__create_fn__'](**locals).__code__
+
+
+def _extract_fields(source):
+    """
+    Replace every named placeholder (like "_field_spam") in the source with a
+    numbered placeholder (like "_field_1"). At the same time, build a mapping
+    of these new numbered placeholders to their actual field names.
+
+    Return the mapping and the patched source.
+    """
+    named_to_numbered = {}  # {"_field_spam": "_field_1", ...}
+    numbered_to_field = {}  # {"_field_1": "spam", ...}
+    parts = _PLACEHOLDER_RE.split(source)
+    # Every other part is a named placeholder. Those are all that we care about:
+    for i in range(1, len(parts), 2):
+        named_placeholder = parts[i]
+        if named_placeholder in named_to_numbered:
+            parts[i] = named_to_numbered[named_placeholder]
+        else:
+            # Using i doesn't give us sequential field numbers, but that's okay:
+            parts[i] = numbered_placeholder = f"{_PLACEHOLDER_PREFIX}{i}"
+            field_name = named_placeholder.removeprefix(_PLACEHOLDER_PREFIX)
+            named_to_numbered[named_placeholder] = numbered_placeholder
+            numbered_to_field[numbered_placeholder] = field_name
+    return numbered_to_field, "".join(parts)
+
+
+def _patch_fields(patches, code):
+    """Copy the given code with its consts, names, and varnames patched."""
+    consts = []
+    for const in code.co_consts:
+        match code.co_name, const:
+            case "__init__", str() if const in patches:
+                # const is a field name:
+                const = patches[const]
+            case "__repr__", str():
+                # const may be a string *containing* field names:
+                unpatched = _PLACEHOLDER_RE.split(const)
+                patched = _apply_patches(patches, unpatched)
+                const = "".join(patched)
+            case "__delattr__" | "__setattr__", tuple():
+                # const may be a tuple *containing* field names:
+                patched = _apply_patches(patches, const)
+                const = tuple(patched)
+        consts.append(const)
+    names = _apply_patches(patches, code.co_names)
+    varnames = _apply_patches(patches, code.co_varnames)
+    return code.replace(
+        co_consts=tuple(consts),
+        co_names=tuple(names),
+        co_varnames=tuple(varnames),
+    )
+
+
+def _apply_patches(patches, unpatched):
+    patched = []
+    for item in unpatched:
+        if item in patches:
+            item = patches[item]
+        patched.append(item)
+    return patched
 
 
 def _field_assign(frozen, name, value, self_name):
@@ -463,19 +570,19 @@ def _field_assign(frozen, name, value, self_name):
     return f'{self_name}.{name}={value}'
 
 
-def _field_init(f, frozen, globals, self_name, slots):
+def _field_init(f, frozen, globals, self_name, slots, i):
     # Return the text of the line in the body of __init__ that will
     # initialize this field.
 
-    default_name = f'_dflt_{f.name}'
+    default_name = f'_dflt_{i}'
     if f.default_factory is not MISSING:
         if f.init:
             # This field has a default factory.  If a parameter is
             # given, use it.  If not, call the factory.
             globals[default_name] = f.default_factory
             value = (f'{default_name}() '
-                     f'if {f.name} is _HAS_DEFAULT_FACTORY '
-                     f'else {f.name}')
+                     f'if {f._placeholder} is _HAS_DEFAULT_FACTORY '
+                     f'else {f._placeholder}')
         else:
             # This is a field that's not in the __init__ params, but
             # has a default factory function.  It needs to be
@@ -498,10 +605,10 @@ def _field_init(f, frozen, globals, self_name, slots):
         if f.init:
             if f.default is MISSING:
                 # There's no default, just do an assignment.
-                value = f.name
+                value = f._placeholder
             elif f.default is not MISSING:
                 globals[default_name] = f.default
-                value = f.name
+                value = f._placeholder
         else:
             # If the class has slots, then initialize this field.
             if slots and f.default is not MISSING:
@@ -520,26 +627,7 @@ def _field_init(f, frozen, globals, self_name, slots):
         return None
 
     # Now, actually generate the field assignment.
-    return _field_assign(frozen, f.name, value, self_name)
-
-
-def _init_param(f):
-    # Return the __init__ parameter string for this field.  For
-    # example, the equivalent of 'x:int=3' (except instead of 'int',
-    # reference a variable set to int, and instead of '3', reference a
-    # variable set to 3).
-    if f.default is MISSING and f.default_factory is MISSING:
-        # There's no default, and no default_factory, just output the
-        # variable name and type.
-        default = ''
-    elif f.default is not MISSING:
-        # There's a default, this will be the name that's used to look
-        # it up.
-        default = f'=_dflt_{f.name}'
-    elif f.default_factory is not MISSING:
-        # There's a factory function.  Set a marker.
-        default = '=_HAS_DEFAULT_FACTORY'
-    return f'{f.name}:_type_{f.name}{default}'
+    return _field_assign(frozen, f._placeholder, value, self_name)
 
 
 def _init_fn(fields, std_fields, kw_only_fields, frozen, has_post_init,
@@ -562,51 +650,76 @@ def _init_fn(fields, std_fields, kw_only_fields, frozen, has_post_init,
                 raise TypeError(f'non-default argument {f.name!r} '
                                 'follows default argument')
 
-    locals = {f'_type_{f.name}': f.type for f in fields}
-    locals.update({
+    locals = {
         'MISSING': MISSING,
         '_HAS_DEFAULT_FACTORY': _HAS_DEFAULT_FACTORY,
         '__dataclass_builtins_object__': object,
-    })
+    }
 
     body_lines = []
-    for f in fields:
-        line = _field_init(f, frozen, locals, self_name, slots)
+    for i, f in enumerate(fields):
+        line = _field_init(f, frozen, locals, self_name, slots, i)
         # line is None means that this field doesn't require
         # initialization (it's a pseudo-field).  Just skip it.
         if line:
             body_lines.append(line)
 
     # Does this class have a post-init function?
     if has_post_init:
-        params_str = ','.join(f.name for f in fields
+        params_str = ','.join(f._placeholder for f in fields
                               if f._field_type is _FIELD_INITVAR)
         body_lines.append(f'{self_name}.{_POST_INIT_NAME}({params_str})')
 
     # If no body lines, use 'pass'.
     if not body_lines:
         body_lines = ['pass']
 
-    _init_params = [_init_param(f) for f in std_fields]
+    _init_params = [f._placeholder for f in std_fields]
     if kw_only_fields:
         # Add the keyword-only args.  Because the * can only be added if
         # there's at least one keyword-only arg, there needs to be a test here
         # (instead of just concatenting the lists together).
         _init_params += ['*']
-        _init_params += [_init_param(f) for f in kw_only_fields]
-    return _create_fn('__init__',
-                      [self_name] + _init_params,
-                      body_lines,
-                      locals=locals,
-                      globals=globals,
-                      return_type=None)
+        _init_params += [f._placeholder for f in kw_only_fields]
+    f = _create_fn('__init__',
+                   (self_name, *_init_params),
+                   tuple(body_lines),
+                   locals=locals,
+                   globals=globals)
+    _add_init_dunders(f, std_fields, kw_only_fields)
+    return f
+
+
+def _add_init_dunders(init, std_fields, kw_only_fields):
+    """
+    Add __annotations__, __defaults__, and __kwdefaults__ to an __init__
+    function.
+    """
+    annotations = {"return": None}
+    defaults = []
+    kwdefaults = {}
+    for field in std_fields:
+        annotations[field.name] = field.type
+        if field.default is not MISSING:
+            defaults.append(field.default)
+        elif field.default_factory is not MISSING:
+            defaults.append(_HAS_DEFAULT_FACTORY)
+    for field in kw_only_fields:
+        annotations[field.name] = field.type
+        if field.default is not MISSING:
+            kwdefaults[field.name] = field.default
+        elif field.default_factory is not MISSING:
+            kwdefaults[field.name] = _HAS_DEFAULT_FACTORY
+    init.__annotations__ = annotations
+    init.__defaults__ = tuple(defaults) or None
+    init.__kwdefaults__ = kwdefaults or None
 
 
 def _repr_fn(fields, globals):
     fn = _create_fn('__repr__',
                     ('self',),
                     ['return self.__class__.__qualname__ + f"(' +
-                     ', '.join([f"{f.name}={{self.{f.name}!r}}"
+                     ', '.join([f"{f._placeholder}={{self.{f._placeholder}!r}}"
                                 for f in fields]) +
                      ')"'],
                      globals=globals)
@@ -618,7 +731,7 @@ def _frozen_get_del_attr(cls, fields, globals):
               'FrozenInstanceError': FrozenInstanceError}
     condition = 'type(self) is cls'
     if fields:
-        condition += ' or name in {' + ', '.join(repr(f.name) for f in fields) + '}'
+        condition += ' or name in {' + ', '.join(repr(f._placeholder) for f in fields) + '}'
     return (_create_fn('__setattr__',
                       ('self', 'name', 'value'),
                       (f'if {condition}:',
@@ -639,8 +752,8 @@ def _frozen_get_del_attr(cls, fields, globals):
 def _cmp_fn(name, op, self_tuple, other_tuple, globals):
     # Create a comparison function.  If the fields in the object are
     # named 'x' and 'y', then self_tuple is the string
-    # '(self.x,self.y)' and other_tuple is the string
-    # '(other.x,other.y)'.
+    # '(self._field_x,self._field_y)' and other_tuple is the string
+    # '(other._field_x,other._field_y)'.
 
     return _create_fn(name,
                       ('self', 'other'),
@@ -756,6 +869,8 @@ def _get_field(cls, a_name, a_type, default_kw_only):
     f.name = a_name
     f.type = a_type
 
+    f._placeholder = f"{_PLACEHOLDER_PREFIX}{a_name}"
+
     # Assume it's a normal field until proven otherwise.  We're next
     # going to decide if it's a ClassVar or InitVar, everything else
     # is just a normal field.