[eudsl-python-extras] remove bytecode dep (#240)

This PR removes our dependency on `bytecode` (which turned out to be
completely superfluous 🤦).

Author: makslevental

projects/eudsl-python-extras/mlir/extras/ast/canonicalize.py

@@ -12,7 +12,6 @@
 from opcode import opmap
 from typing import List, Union, Sequence, get_type_hints
 
-from bytecode import ConcreteBytecode
 
 from ..ast.util import get_module_cst, set_lineno, find_func_in_code_object
 
@@ -166,25 +165,23 @@ def to_str(self: OpCode):
 setattr(OpCode, "__str__", to_str)
 
 
-class BytecodePatcher(ABC):
+class FunctionPatcher(ABC):
     def __init__(self, context=None):
         self.context = context
 
     @abstractmethod
-    def patch_bytecode(self, code: ConcreteBytecode, original_f) -> ConcreteBytecode:
+    def patch_function(self, original_f):
         pass
 
 
-def patch_bytecode(f, patchers: List[type(BytecodePatcher)] = None):
+def patch_function(f, patchers: List[type(FunctionPatcher)] = None):
     if patchers is None:
         return f
-    code = ConcreteBytecode.from_code(f.__code__)
     context = types.SimpleNamespace()
     for patcher in patchers:
-        code = patcher(context).patch_bytecode(code, f)
+        new_f = patcher(context).patch_function(f)
 
-    f.__code__ = code.to_code()
-    return f
+    return new_f
 
 
 class Canonicalizer(ABC):
@@ -195,22 +192,22 @@ def cst_transformers(self) -> List[StrictTransformer]:
 
     @property
     @abstractmethod
-    def bytecode_patchers(self) -> List[BytecodePatcher]:
+    def function_patchers(self) -> List[FunctionPatcher]:
         pass
 
 
 def canonicalize(*, using: Union[Canonicalizer, Sequence[Canonicalizer]]):
     if not isinstance(using, Sequence):
         using = [using]
     cst_transformers = []
-    bytecode_patchers = []
+    function_patchers = []
     for u in using:
         cst_transformers.extend(u.cst_transformers)
-        bytecode_patchers.extend(u.bytecode_patchers)
+        function_patchers.extend(u.function_patchers)
 
     def wrapper(f):
         f = transform_ast(f, cst_transformers)
-        f = patch_bytecode(f, bytecode_patchers)
+        f = patch_function(f, function_patchers)
         return f
 
     return wrapper

projects/eudsl-python-extras/mlir/extras/ast/util.py

@@ -6,11 +6,9 @@
 import inspect
 import io
 import types
-from opcode import opmap
 from textwrap import dedent
 from typing import Dict
 
-from bytecode import ConcreteBytecode
 from cloudpickle import cloudpickle
 
 from ...ir import Type
@@ -87,39 +85,6 @@ def make_cell(value=_empty_cell_value):
     return cell
 
 
-# based on https://github.com/python/cpython/blob/a4b44d39cd6941cc03590fee7538776728bdfd0a/Lib/test/test_code.py#L197
-def replace_closure(code, new_closure: Dict):
-    COPY_FREE_VARS = opmap["COPY_FREE_VARS"]
-    LOAD_DEREF = opmap["LOAD_DEREF"]
-
-    # get the orig localplus that will be loaded from by the orig bytecode LOAD_DEREF arg_i
-    localsplus, _localsplus_name_to_idx = get_localsplus_name_to_idx(code)
-
-    # closure vars go into co_freevars
-    new_code = code.replace(co_freevars=tuple(new_closure.keys()))
-    # closure is a tuple of cells
-    closure = tuple(
-        make_cell(v) if not isinstance(v, types.CellType) else v
-        for v in new_closure.values()
-    )
-
-    new_code = ConcreteBytecode.from_code(new_code)
-    # update how many closure vars are loaded from frame
-    # see https://github.com/python/cpython/blob/6078f2033ea15a16cf52fe8d644a95a3be72d2e3/Python/bytecodes.c#L1571
-    assert new_code[0].opcode == COPY_FREE_VARS, f"{new_code[0].opcode=}"
-    new_code[0].arg = len(closure)
-
-    # map orig localsplus arg_i to new localplus position/arg_i
-    new_localsplus = new_code.varnames + new_code.cellvars + new_code.freevars
-    new_localsplus_name_to_idx = {v: i for i, v in enumerate(new_localsplus)}
-    for c in new_code:
-        if c.opcode == LOAD_DEREF and c.arg < len(localsplus):
-            c.arg = new_localsplus_name_to_idx[localsplus[c.arg]]
-    new_code = new_code.to_code()
-
-    return new_code, closure
-
-
 def unpickle_mlir_type(v):
     return Type.parse(v)
 
@@ -145,7 +110,13 @@ def copy_object(obj):
 # potentially more complete approach https://stackoverflow.com/a/56901529/9045206
 def copy_func(f, new_closure: Dict = None):
     if new_closure is not None:
-        code, closure = replace_closure(f.__code__, new_closure)
+        # closure vars go into co_freevars
+        code = f.__code__.replace(co_freevars=tuple(new_closure.keys()))
+        # closure is a tuple of cells
+        closure = tuple(
+            make_cell(v) if not isinstance(v, types.CellType) else v
+            for v in new_closure.values()
+        )
     else:
         closure = copy_object(f.__closure__)
         code = f.__code__

projects/eudsl-python-extras/mlir/extras/dialects/arith.py

@@ -10,9 +10,8 @@
 from typing import Optional, Tuple, Union
 
 import numpy as np
-from bytecode import ConcreteBytecode
 
-from ..ast.canonicalize import StrictTransformer, Canonicalizer, BytecodePatcher
+from ..ast.canonicalize import StrictTransformer, Canonicalizer, FunctionPatcher
 from ..util import infer_mlir_type, mlir_type_to_np_dtype
 from ..._mlir_libs._mlir import register_value_caster
 from ...dialects import complex as complex_dialect
@@ -596,19 +595,19 @@ def visit_AugAssign(
         return updated_node
 
 
-class ArithPatchByteCode(BytecodePatcher):
-    def patch_bytecode(self, code: ConcreteBytecode, f):
+class ArithPatchFunction(FunctionPatcher):
+    def patch_function(self, f):
         # TODO(max): this is bad and should be in the closure rather than as a global
         from ...dialects import arith, math
 
         f.__globals__["math_dialect"] = math
         f.__globals__["arith_dialect"] = arith
-        return code
+        return f
 
 
 class ArithCanonicalizer(Canonicalizer):
     cst_transformers = [CanonicalizeFMA]
-    bytecode_patchers = [ArithPatchByteCode]
+    function_patchers = [ArithPatchFunction]
 
 
 canonicalizer = ArithCanonicalizer()

projects/eudsl-python-extras/mlir/extras/dialects/scf.py

@@ -7,13 +7,11 @@
 from copy import deepcopy
 from typing import List, Union, Optional, Sequence
 
-from bytecode import ConcreteBytecode
-
 from .arith import constant as _ext_arith_constant, index_cast
-from ..ast.canonicalize import BytecodePatcher, Canonicalizer, StrictTransformer
+from ..ast.canonicalize import FunctionPatcher, Canonicalizer, StrictTransformer
 from ..ast.util import ast_call, set_lineno, append_hidden_node
 from ..meta import region_op
-from ..util import get_user_code_loc, region_adder
+from ..util import region_adder
 from ...dialects._ods_common import (
     _cext,
     get_default_loc_context,
@@ -567,14 +565,14 @@ def visit_If(self, updated_node: ast.If) -> Union[ast.With, List[ast.With]]:
             return then_with
 
 
-class RemoveJumpsAndInsertGlobals(BytecodePatcher):
-    def patch_bytecode(self, code: ConcreteBytecode, f):
+class RemoveJumpsAndInsertGlobals(FunctionPatcher):
+    def patch_function(self, f):
         # TODO(max): this is bad and should be in the closure rather than as a global
         f.__globals__[yield_.__name__] = yield_
         f.__globals__[if_ctx_manager.__name__] = if_ctx_manager
         f.__globals__[else_ctx_manager.__name__] = else_ctx_manager
         f.__globals__[placeholder_opaque_t.__name__] = placeholder_opaque_t
-        return code
+        return f
 
 
 class SCFCanonicalizer(Canonicalizer):
@@ -586,7 +584,7 @@ class SCFCanonicalizer(Canonicalizer):
         CanonicalizeWhile,
     ]
 
-    bytecode_patchers = [RemoveJumpsAndInsertGlobals]
+    function_patchers = [RemoveJumpsAndInsertGlobals]
 
 
 canonicalizer = SCFCanonicalizer()

projects/eudsl-python-extras/requirements.txt

@@ -1,4 +1,3 @@
 PyYAML>=5.4.0
-bytecode>=0.17.0
 cloudpickle>=3.0.0
 numpy>=1.19.5, <=2.1.2