retrofy 0.2.1__py3-none-any.whl

retrofy/__init__.py

@@ -0,0 +1,8 @@
+"""
+Documentation for the retrofy package
+
+"""
+
+# Import the version from the generated _version.py file. __version__ is part
+# of the public API, and we therefore ignore the "unused" (F401) lint warning.
+from ._version import __version__  # noqa: F401  # pylint: disable=import-error

retrofy/_converters.py

@@ -0,0 +1,174 @@
+import dataclasses
+import typing
+
+import libcst as cst
+
+from ._transformations import dataclass, match_statement, type_alias, walrus
+
+
+class TypingTransformer(cst.CSTTransformer):
+    def __init__(self, scope):
+        self._scope = scope
+        self._require_typing = False
+        self._has_typing = False  # TODO: We can figure this out.
+
+    def leave_Annotation(
+        self,
+        node: cst.Annotation,
+        updated_node: cst.Annotation,
+    ) -> cst.Annotation:
+        if isinstance(updated_node.annotation, cst.BinaryOperation):
+            # TODO: Use a matcher here.
+            if isinstance(updated_node.annotation.operator, cst.BitOr):
+                # print('ANNO scope', self._scope[node.annotation])
+                self._require_typing = True
+                new_node = cst.Subscript(
+                    cst.Attribute(
+                        cst.Name("typing"),
+                        cst.Name("Union"),
+                    ),
+                    slice=(
+                        cst.SubscriptElement(
+                            updated_node.annotation.left,  # type: ignore
+                        ),
+                        cst.SubscriptElement(
+                            updated_node.annotation.right,  # type: ignore
+                        ),
+                    ),
+                )
+                return cst.Annotation(new_node)
+
+        return updated_node
+
+    def leave_Module(self, node: cst.Module, updated_node: cst.Module) -> cst.Module:
+        if self._require_typing:
+            # Find the correct position to insert the typing import
+            # It should be after any docstrings and __future__ imports
+            insert_position = self._find_typing_import_position(updated_node.body)
+
+            typing_import = cst.SimpleStatementLine(
+                [
+                    cst.Import(
+                        [
+                            cst.ImportAlias(
+                                cst.Name("typing"),
+                            ),
+                        ],
+                    ),
+                ],
+                trailing_whitespace=cst.TrailingWhitespace(
+                    newline=cst.Newline(),
+                ),
+            )
+
+            new_stmts = list(updated_node.body)
+            new_stmts.insert(insert_position, typing_import)
+
+            return dataclasses.replace(updated_node, body=tuple(new_stmts))
+        return updated_node
+
+    def _find_typing_import_position(
+        self,
+        body: typing.Tuple[cst.BaseStatement, ...],
+    ) -> int:
+        """Find the correct position to insert the typing import."""
+        position = 0
+
+        # Skip module docstrings
+        if body and isinstance(body[0], cst.SimpleStatementLine):
+            if (
+                len(body[0].body) == 1
+                and isinstance(body[0].body[0], cst.Expr)
+                and isinstance(body[0].body[0].value, cst.SimpleString)
+            ):
+                position = 1
+
+        # Skip __future__ imports
+        for i in range(position, len(body)):
+            stmt = body[i]
+            if isinstance(stmt, cst.SimpleStatementLine):
+                for substmt in stmt.body:
+                    if (
+                        isinstance(substmt, cst.ImportFrom)
+                        and substmt.module
+                        and isinstance(substmt.module, cst.Attribute)
+                        and substmt.module.attr.value == "__future__"
+                    ):
+                        position = i + 1
+                        break
+                    elif (
+                        isinstance(substmt, cst.ImportFrom)
+                        and substmt.module
+                        and isinstance(substmt.module, cst.Name)
+                        and substmt.module.value == "__future__"
+                    ):
+                        position = i + 1
+                        break
+                else:
+                    # If we didn't find a __future__ import in this statement, stop looking
+                    break
+            else:
+                # If we hit a non-simple statement, stop looking
+                break
+
+        return position
+
+
+def convert_union(module: cst.Module) -> cst.Module:
+    """
+    Given typing such as `SomeClass | AnotherClass`, convert this to
+    `typing.Union[SomeClass, AnotherClass]`, with the appropriate `typing`
+    import included.
+
+    Note that this typically doesn't need to be done at runtime, and mostly is
+    a type checking feature.
+
+    """
+
+    wrapper = cst.metadata.MetadataWrapper(module)
+    metadata = wrapper.resolve(cst.metadata.ScopeProvider)
+
+    transformer = TypingTransformer(metadata)
+    return module.visit(transformer)
+
+
+def convert_sequence_subscript(module: cst.Module) -> cst.Module:
+    """
+    Convert the built-in sequence types such as `list[int]` to the
+    `typing.List[int]` form.
+
+    Currently only `list[int]` is supported.
+    """
+    # TODO: Do this correctly. Including adding the necessry typing import.
+    orig = module.code
+    new = orig.replace("list[str]", "typing.List[str]")
+    if new != orig:
+        return cst.parse_module(new)
+    return module
+
+
+def convert_walrus_operator(module: cst.Module) -> cst.Module:
+    return module.visit(walrus.WalrusOperatorTransformer())
+
+
+def convert_type_alias(module: cst.Module) -> cst.Module:
+    return module.visit(type_alias.PEP695Transformer())
+
+
+def convert_match_statement(module: cst.Module) -> cst.Module:
+    return module.visit(match_statement.MatchStatementTransformer())
+
+
+def convert_dataclass(module: cst.Module) -> cst.Module:
+    return module.visit(dataclass.DataclassTransformer())
+
+
+def convert(code: str) -> str:
+    mod = cst.parse_module(code)
+    mod = convert_sequence_subscript(mod)
+    mod = convert_walrus_operator(mod)
+    mod = convert_type_alias(mod)
+    mod = convert_dataclass(mod)
+    mod = convert_match_statement(mod)
+    mod = convert_union(mod)
+    return mod.code

retrofy/_meta_hook_converter.py

@@ -0,0 +1,62 @@
+import sys
+import typing
+from importlib.abc import MetaPathFinder, SourceLoader
+
+from ._converters import convert
+
+
+class OnTheFlyConverter(SourceLoader):
+    def __init__(self, path):
+        self.path = path
+
+    def get_filename(self, fullname):
+        return self.path
+
+    def get_data(self, filename):
+        """exec_module is already defined for us, we just have to provide a way
+        of getting the source code of the module"""
+        with open(filename) as f:
+            data = f.read()
+
+        new_code = convert(data)
+        return new_code
+
+
+class MyMetaPathFinder(MetaPathFinder):
+    def __init__(self, package_names: typing.Sequence[str] = ()):
+        self.package_names: set[str] = set()
+        self.add_package_handling(package_names)
+
+    def add_package_handling(self, package_names: typing.Sequence[str]):
+        self.package_names.update(package_names)
+
+    def find_spec(self, fullname, path, target=None):
+        for prefix in self.package_names:
+            if fullname != prefix and not fullname.startswith(f"{prefix}."):
+                return None
+
+        # If your custom logic doesn't handle it, defer to the next finder
+        for finder in sys.meta_path:
+            if isinstance(finder, MyMetaPathFinder):
+                continue
+
+            spec = finder.find_spec(fullname, path, target)
+            if spec:
+                break
+        else:
+            return None
+
+        spec.loader = OnTheFlyConverter(spec.origin)
+        return spec
+
+
+def register_hook(package_names):
+    for finder in sys.meta_path:
+        if isinstance(finder, MyMetaPathFinder):
+            existing_hook = finder
+            break
+    else:
+        existing_hook = MyMetaPathFinder()
+        sys.meta_path.insert(0, MyMetaPathFinder(package_names))
+
+    existing_hook.add_package_handling(package_names)

retrofy/_pep517_hooks.py

@@ -0,0 +1,71 @@
+import pathlib
+import shutil
+import zipfile
+
+from setuptools_ext import WheelModifier
+
+from ._converters import convert
+
+
+def compatibility_via_import_hook(wheel: pathlib.Path):
+    """
+    Add a pth hook to ensure make imported code compatible at import time
+    (i.e. suitable for editable mode)
+
+    """
+    editable_copy = wheel.parent / (wheel.name + ".copy.whl")
+    shutil.copy(wheel, editable_copy)
+
+    with zipfile.ZipFile(str(editable_copy), "r") as whl_zip:
+        whl = WheelModifier(whl_zip)
+
+        top_level = (
+            whl.read(
+                whl.dist_info_dirname() + "/top_level.txt",
+            )
+            .decode("utf-8")
+            .splitlines()
+        )
+        top_level_pkgs = [pkg for pkg in top_level if pkg]
+
+        for pkg in top_level_pkgs:
+            fn = f"_retrofy.__editable_compat__.{pkg}.pth"
+            script = (
+                f"import retrofy._meta_hook_converter as c; c.register_hook(['{pkg}']);"
+            )
+            whl.write(zipfile.ZipInfo(fn), script)
+
+        with wheel.open("wb") as whl_fh:
+            whl.write_wheel(whl_fh)
+
+    print("Enabling automatic retrofiting of Python code at import-time")
+
+    # TODO: I think we need to add retrofy as a runtime dependency for editable mode
+    #  (we can do this).
+
+    editable_copy.unlink()
+
+
+def compatibility_via_rewrite(wheel: pathlib.Path):
+    """Change code within the given wheel to be compatible"""
+    editable_copy = wheel.parent / (wheel.name + ".copy.whl")
+    shutil.copy(wheel, editable_copy)
+
+    has_modifications = False
+
+    with zipfile.ZipFile(str(editable_copy), "r") as whl_zip:
+        whl = WheelModifier(whl_zip)
+
+        for filename in whl_zip.namelist():
+            if filename.endswith(".py"):
+                code = whl.read(filename).decode("utf-8")
+                new_code = convert(code)
+                if new_code != code:
+                    print(f"Converted {filename} to compatibility syntax")
+                    whl.write(filename, new_code)
+                    has_modifications = True
+        if has_modifications:
+            with wheel.open("wb") as whl_fh:
+                whl.write_wheel(whl_fh)
+
+    editable_copy.unlink()

retrofy/_transformations/dataclass.py

@@ -0,0 +1,240 @@
+from __future__ import annotations
+
+from typing import List
+
+import libcst as cst
+
+
+class DataclassTransformer(cst.CSTTransformer):
+    """
+    Transform @dataclass decorated classes to include __match_args__ attribute.
+
+    This transformation adds __match_args__ to dataclasses for Python < 3.10
+    compatibility, unless match_args=False is explicitly specified.
+    """
+
+    def leave_Module(
+        self,
+        original_node: cst.Module,
+        updated_node: cst.Module,
+    ) -> cst.Module:
+        """Transform the entire module to add __match_args__ handling after class definitions."""
+        new_body = []
+
+        for stmt in updated_node.body:
+            if isinstance(stmt, cst.ClassDef):
+                # Process the class
+                processed_class, needs_deletion = self._process_dataclass(stmt)
+                new_body.append(processed_class)
+
+                # Add deletion statement if needed
+                if needs_deletion:
+                    delete_stmt = self._create_delete_statement(
+                        processed_class.name.value,
+                    )
+                    # Add the delete statement with proper spacing
+                    delete_stmt = delete_stmt.with_changes(
+                        leading_lines=[
+                            cst.EmptyLine(
+                                whitespace=cst.SimpleWhitespace(""),
+                                comment=None,
+                            ),
+                        ],
+                    )
+                    new_body.append(delete_stmt)
+            else:
+                new_body.append(stmt)
+
+        return updated_node.with_changes(body=new_body)
+
+    def _process_dataclass(self, class_def: cst.ClassDef) -> tuple[cst.ClassDef, bool]:
+        """Process a class definition and return (modified_class, needs_deletion)."""
+
+        # Check if this class has a @dataclass decorator
+        if not self._has_dataclass_decorator(class_def):
+            return class_def, False
+
+        # Check if match_args=False is specified
+        if self._has_match_args_false(class_def):
+            # Remove match_args=False parameter and mark for deletion
+            modified_class = self._remove_match_args_parameter(class_def)
+            # Only need deletion if __match_args__ wasn't explicitly defined
+            needs_deletion = not self._has_match_args_attribute(class_def)
+            return modified_class, needs_deletion
+
+        # Check if __match_args__ is already defined
+        if self._has_match_args_attribute(class_def):
+            return class_def, False
+
+        # Add __match_args__ for regular dataclasses
+        field_names = self._extract_field_names(class_def)
+        if field_names:
+            match_args_stmt = self._create_match_args_statement(field_names)
+            new_body = list(class_def.body.body) + [match_args_stmt]
+            new_body_node = class_def.body.with_changes(body=new_body)
+            modified_class = class_def.with_changes(body=new_body_node)
+            return modified_class, False
+
+        return class_def, False
+
+    def _create_delete_statement(self, class_name: str) -> cst.Try:
+        """Create try/except block to delete __match_args__ from the specified class."""
+        # try: del ClassName.__match_args__; except AttributeError: pass
+
+        del_stmt = cst.Del(
+            target=cst.Attribute(
+                value=cst.Name(class_name),
+                attr=cst.Name("__match_args__"),
+            ),
+        )
+
+        try_body = cst.IndentedBlock([cst.SimpleStatementLine([del_stmt])])
+
+        except_handler = cst.ExceptHandler(
+            type=cst.Name("AttributeError"),
+            body=cst.IndentedBlock([cst.SimpleStatementLine([cst.Pass()])]),
+        )
+
+        return cst.Try(
+            body=try_body,
+            handlers=[except_handler],
+        )
+
+    def _has_dataclass_decorator(self, class_def: cst.ClassDef) -> bool:
+        """Check if class has @dataclass decorator."""
+        for decorator in class_def.decorators:
+            if isinstance(decorator.decorator, cst.Name):
+                if decorator.decorator.value == "dataclass":
+                    return True
+            elif isinstance(decorator.decorator, cst.Call):
+                if (
+                    isinstance(decorator.decorator.func, cst.Name)
+                    and decorator.decorator.func.value == "dataclass"
+                ):
+                    return True
+        return False
+
+    def _has_match_args_false(self, class_def: cst.ClassDef) -> bool:
+        """Check if @dataclass has match_args=False."""
+        for decorator in class_def.decorators:
+            if isinstance(decorator.decorator, cst.Call):
+                if (
+                    isinstance(decorator.decorator.func, cst.Name)
+                    and decorator.decorator.func.value == "dataclass"
+                ):
+                    # Check for match_args=False in arguments
+                    for arg in decorator.decorator.args:
+                        if (
+                            isinstance(arg, cst.Arg)
+                            and isinstance(arg.keyword, cst.Name)
+                            and arg.keyword.value == "match_args"
+                            and isinstance(arg.value, cst.Name)
+                            and arg.value.value == "False"
+                        ):
+                            return True
+        return False
+
+    def _has_match_args_attribute(self, class_def: cst.ClassDef) -> bool:
+        """Check if class already has __match_args__ attribute."""
+        for stmt in class_def.body.body:
+            if isinstance(stmt, cst.SimpleStatementLine):
+                for inner_stmt in stmt.body:
+                    if isinstance(inner_stmt, cst.Assign):
+                        for target in inner_stmt.targets:
+                            if (
+                                isinstance(target.target, cst.Name)
+                                and target.target.value == "__match_args__"
+                            ):
+                                return True
+            elif isinstance(stmt, cst.AnnAssign):
+                if (
+                    isinstance(stmt.target, cst.Name)
+                    and stmt.target.value == "__match_args__"
+                ):
+                    return True
+        return False
+
+    def _extract_field_names(self, class_def: cst.ClassDef) -> List[str]:
+        """Extract field names from dataclass."""
+        field_names = []
+
+        for stmt in class_def.body.body:
+            if isinstance(stmt, cst.AnnAssign):
+                # Annotated assignment: x: int or x: int = 5
+                if isinstance(stmt.target, cst.Name):
+                    field_names.append(stmt.target.value)
+            elif isinstance(stmt, cst.SimpleStatementLine):
+                for inner_stmt in stmt.body:
+                    if isinstance(inner_stmt, cst.Assign):
+                        # Regular assignment: x = 5 (less common in dataclasses)
+                        for target in inner_stmt.targets:
+                            if isinstance(target.target, cst.Name):
+                                field_names.append(target.target.value)
+                    elif isinstance(inner_stmt, cst.AnnAssign):
+                        # Annotated assignment in simple statement line
+                        if isinstance(inner_stmt.target, cst.Name):
+                            field_names.append(inner_stmt.target.value)
+
+        return field_names
+
+    def _create_match_args_statement(
+        self,
+        field_names: List[str],
+    ) -> cst.SimpleStatementLine:
+        """Create __match_args__ = ('field1', 'field2', ...) statement."""
+
+        # Create tuple elements from field names
+        elements = []
+        for name in field_names:
+            elements.append(cst.Element(cst.SimpleString(f"'{name}'")))
+
+        # Create tuple
+        if len(elements) == 1:
+            # Single element tuple needs trailing comma
+            tuple_value = cst.Tuple([elements[0].with_changes(comma=cst.Comma())])
+        else:
+            tuple_value = cst.Tuple(elements)
+
+        # Create assignment: __match_args__ = (...)
+        assignment = cst.Assign(
+            targets=[cst.AssignTarget(cst.Name("__match_args__"))],
+            value=tuple_value,
+        )
+
+        return cst.SimpleStatementLine([assignment])
+
+    def _remove_match_args_parameter(self, class_def: cst.ClassDef) -> cst.ClassDef:
+        """Remove match_args=False parameter from @dataclass decorator."""
+        new_decorators = []
+
+        for decorator in class_def.decorators:
+            if isinstance(decorator.decorator, cst.Call):
+                if (
+                    isinstance(decorator.decorator.func, cst.Name)
+                    and decorator.decorator.func.value == "dataclass"
+                ):
+                    # Filter out the match_args parameter
+                    new_args = []
+                    for arg in decorator.decorator.args:
+                        if not (
+                            isinstance(arg, cst.Arg)
+                            and isinstance(arg.keyword, cst.Name)
+                            and arg.keyword.value == "match_args"
+                        ):
+                            new_args.append(arg)
+
+                    # If no args remain, convert back to simple @dataclass
+                    if not new_args:
+                        new_decorator = cst.Decorator(cst.Name("dataclass"))
+                    else:
+                        new_decorator = decorator.with_changes(
+                            decorator=decorator.decorator.with_changes(args=new_args),
+                        )
+
+                    new_decorators.append(new_decorator)
+                else:
+                    new_decorators.append(decorator)
+            else:
+                new_decorators.append(decorator)
+
+        return class_def.with_changes(decorators=new_decorators)

retrofy/_transformations/import_utils.py

@@ -0,0 +1,103 @@
+"""Utilities for managing imports in transformations."""
+
+from typing import Tuple
+
+import libcst as cst
+
+
+class ImportManager:
+    """Helper class for managing automatic imports in transformations."""
+
+    def __init__(self):
+        self._required_imports = set()
+
+    def require_import(self, import_name: str) -> None:
+        """Mark an import as required."""
+        self._required_imports.add(import_name)
+
+    def apply_imports(self, module: cst.Module) -> cst.Module:
+        """Add all required imports to the module."""
+        if not self._required_imports:
+            return module
+
+        # Find the correct position to insert imports
+        insert_position = self._find_import_position(module.body)
+
+        new_stmts = list(module.body)
+        for import_name in sorted(self._required_imports):
+            import_stmt = self._create_import_statement(import_name)
+            new_stmts.insert(insert_position, import_stmt)
+            insert_position += 1
+
+        return module.with_changes(body=tuple(new_stmts))
+
+    def _create_import_statement(self, import_name: str) -> cst.SimpleStatementLine:
+        """Create an import statement for the given import name."""
+        if "." in import_name:
+            # Handle dotted imports like "collections.abc"
+            parts = import_name.split(".")
+            import_node = cst.Name(parts[0])
+            for part in parts[1:]:
+                import_node = cst.Attribute(import_node, cst.Name(part))
+        else:
+            # Simple import like "typing"
+            import_node = cst.Name(import_name)
+
+        return cst.SimpleStatementLine(
+            [
+                cst.Import(
+                    [
+                        cst.ImportAlias(import_node),
+                    ],
+                ),
+            ],
+            trailing_whitespace=cst.TrailingWhitespace(
+                newline=cst.Newline(),
+            ),
+        )
+
+    def _find_import_position(
+        self,
+        body: Tuple[cst.BaseStatement, ...],
+    ) -> int:
+        """Find the correct position to insert imports."""
+        position = 0
+
+        # Skip module docstrings
+        if body and isinstance(body[0], cst.SimpleStatementLine):
+            if (
+                len(body[0].body) == 1
+                and isinstance(body[0].body[0], cst.Expr)
+                and isinstance(body[0].body[0].value, cst.SimpleString)
+            ):
+                position = 1
+
+        # Skip __future__ imports
+        for i in range(position, len(body)):
+            stmt = body[i]
+            if isinstance(stmt, cst.SimpleStatementLine):
+                for substmt in stmt.body:
+                    if (
+                        isinstance(substmt, cst.ImportFrom)
+                        and substmt.module
+                        and isinstance(substmt.module, cst.Attribute)
+                        and substmt.module.attr.value == "__future__"
+                    ):
+                        position = i + 1
+                        break
+                    elif (
+                        isinstance(substmt, cst.ImportFrom)
+                        and substmt.module
+                        and isinstance(substmt.module, cst.Name)
+                        and substmt.module.value == "__future__"
+                    ):
+                        position = i + 1
+                        break
+                else:
+                    # If we didn't find a __future__ import in this statement, stop looking
+                    break
+            else:
+                # If we hit a non-simple statement, stop looking
+                break
+
+        return position