<commit_before>import ast
import inspect
import cgen as c


class IntrinsicNode(ast.AST):
    def __init__(self, obj, ccode):
        self.obj = obj
        self.ccode = ccode


class GridNode(IntrinsicNode):
    def __getattr__(self, attr):
        return FieldNode(getattr(self.obj, attr),
                         ccode="%s->%s" % (self.ccode, attr))


class FieldNode(IntrinsicNode):
    def __getitem__(self, attr):
        return IntrinsicNode(None, ccode=self.obj.ccode_subscript(*attr))


class ParticleAttributeNode(IntrinsicNode):
    def __init__(self, obj, attr):
        self.obj = obj
        self.attr = attr
        self.ccode = "%s->%s" % (obj.ccode, attr)
        self.ccode_index_var = None

        if self.attr == 'lon':
            self.ccode_index_var = "%s->%s" % (self.obj.ccode, "xi")
        elif self.attr == 'lat':
            self.ccode_index_var = "%s->%s" % (self.obj.ccode, "yi")

    @property
    def pyast_index_update(self):
        pyast = ast.Assign()
        pyast.targets = [IntrinsicNode(None, ccode=self.ccode_index_var)]
        pyast.value = IntrinsicNode(None, ccode=self.ccode_index_update)
        return pyast

    @property
    def ccode_index_update(self):
        """C-code for the index update requires after updating p.lon/p.lat"""
        if self.attr == 'lon':
            return "advance_index(%s, %s, GRID_XDIM, U_lon)" \
                % (self.ccode, self.ccode_index_var)
        if self.attr == 'lat':
            return "advance_index(%s, %s, GRID_YDIM, U_lat)" \
                % (self.ccode, self.ccode_index_var)
        return ""


class ParticleNode(IntrinsicNode):
    def __getattr__(self, attr):
        if attr in self.obj.var_types:
            return ParticleAttributeNode(self, attr)
        else:
            raise AttributeError("""Particle type %s does not define attribute "%s".
Please add '%s' to %s.users_vars or define an appropriate sub-class."""
                                 % (self.obj, attr, attr, self.obj))


class IntrinsicTransformer(ast.NodeTransformer):
    """AST transformer that catches any mention of intrinsic variable
    names, such as 'particle' or 'grid', inserts placeholder objects
    and propagates attribute access."""

    def __init__(self, grid, ptype):
        self.grid = grid
        self.ptype = ptype

    def visit_Name(self, node):
        if node.id == 'grid':
            return GridNode(self.grid, ccode='grid')
        elif node.id == 'particle':
            return ParticleNode(self.ptype, ccode='particle')
        else:
            return node

    def visit_Attribute(self, node):
        node.value = self.visit(node.value)
        if isinstance(node.value, IntrinsicNode):
            return getattr(node.value, node.attr)
        else:
            raise NotImplementedError("Cannot propagate attribute access to C-code")

    def visit_AugAssign(self, node):
        node.target = self.visit(node.target)
        node.op = self.visit(node.op)
        node.value = self.visit(node.value)

        # Capture p.lat/p.lon updates and insert p.xi/p.yi updates
        if isinstance(node.target, ParticleAttributeNode) \
           and node.target.ccode_index_var is not None:
            node = [node, node.target.pyast_index_update]
        return node

    def visit_Assign(self, node):
        node.targets = [self.visit(t) for t in node.targets]
        node.value = self.visit(node.value)

        # Capture p.lat/p.lon updates and insert p.xi/p.yi updates
        if isinstance(node.targets[0], ParticleAttributeNode) \
           and node.targets[0].ccode_index_var is not None:
            node = [node, node.targets[0].pyast_index_update]
        return node


class TupleSplitter(ast.NodeTransformer):
    """AST transformer that detects and splits Pythonic tuple
    assignments into multiple statements for conversion to C."""

    def visit_Assign(self, node):
        if isinstance(node.targets[0], ast.Tuple) \
           and isinstance(node.value, ast.Tuple):
            t_elts = node.targets[0].elts
            v_elts = node.value.elts
            if len(t_elts) != len(v_elts):
                raise AttributeError("Tuple lenghts in assignment do not agree")
            node = [ast.Assign() for _ in t_elts]
            for n, t, v in zip(node, t_elts, v_elts):
                n.targets = [t]
                n.value = v
        return node


class KernelGenerator(ast.NodeVisitor):
    """Code generator class that translates simple Python kernel
    functions into C functions by populating and accessing the `ccode`
    attriibute on nodes in the Python AST."""

    # Intrinsic variables that appear as function arguments
    kernel_vars = ['particle', 'grid', 'dt']

    def __init__(self, grid, ptype):
        self.grid = grid
        self.ptype = ptype

    def generate(self, pyfunc):
        # Parse the Python code into an AST
        self.py_ast = ast.parse(inspect.getsource(pyfunc.func_code))

        # Untangle Pythonic tuple-assignment statements
        self.py_ast = TupleSplitter().visit(self.py_ast)

        # Replace occurences of intrinsic objects in Python AST
        transformer = IntrinsicTransformer(self.grid, self.ptype)
        self.py_ast = transformer.visit(self.py_ast.body[0])

        # Generate C-code for all nodes in the Python AST
        self.visit(self.py_ast)
        self.ccode = self.py_ast.ccode

        # Derive local function variables and insert declaration
        funcvars = list(pyfunc.func_code.co_varnames)
        for kvar in self.kernel_vars:
            funcvars.remove(kvar)
        self.ccode.body.insert(0, c.Value("float", ", ".join(funcvars)))

        return self.ccode

    def visit_FunctionDef(self, node):
        # Create function declaration and argument list
        decl = c.Static(c.DeclSpecifier(c.Value("void", node.name), spec='inline'))
        U, V = (self.grid.U, self.grid.V)
        args = [c.Pointer(c.Value(self.ptype.name, "particle")), c.Value("float", "dt"),
                c.ArrayOf(c.Value("float", U.ccode_lon), count=U.lon.size),
                c.ArrayOf(c.Value("float", U.ccode_lat), count=U.lat.size),<commit_msg>Python3: Change `func_code` to `__code__`<commit_after>import ast
import inspect
import cgen as c


class IntrinsicNode(ast.AST):
    def __init__(self, obj, ccode):
        self.obj = obj
        self.ccode = ccode


class GridNode(IntrinsicNode):
    def __getattr__(self, attr):
        return FieldNode(getattr(self.obj, attr),
                         ccode="%s->%s" % (self.ccode, attr))


class FieldNode(IntrinsicNode):
    def __getitem__(self, attr):
        return IntrinsicNode(None, ccode=self.obj.ccode_subscript(*attr))


class ParticleAttributeNode(IntrinsicNode):
    def __init__(self, obj, attr):
        self.obj = obj
        self.attr = attr
        self.ccode = "%s->%s" % (obj.ccode, attr)
        self.ccode_index_var = None

        if self.attr == 'lon':
            self.ccode_index_var = "%s->%s" % (self.obj.ccode, "xi")
        elif self.attr == 'lat':
            self.ccode_index_var = "%s->%s" % (self.obj.ccode, "yi")

    @property
    def pyast_index_update(self):
        pyast = ast.Assign()
        pyast.targets = [IntrinsicNode(None, ccode=self.ccode_index_var)]
        pyast.value = IntrinsicNode(None, ccode=self.ccode_index_update)
        return pyast

    @property
    def ccode_index_update(self):
        """C-code for the index update requires after updating p.lon/p.lat"""
        if self.attr == 'lon':
            return "advance_index(%s, %s, GRID_XDIM, U_lon)" \
                % (self.ccode, self.ccode_index_var)
        if self.attr == 'lat':
            return "advance_index(%s, %s, GRID_YDIM, U_lat)" \
                % (self.ccode, self.ccode_index_var)
        return ""


class ParticleNode(IntrinsicNode):
    def __getattr__(self, attr):
        if attr in self.obj.var_types:
            return ParticleAttributeNode(self, attr)
        else:
            raise AttributeError("""Particle type %s does not define attribute "%s".
Please add '%s' to %s.users_vars or define an appropriate sub-class."""
                                 % (self.obj, attr, attr, self.obj))


class IntrinsicTransformer(ast.NodeTransformer):
    """AST transformer that catches any mention of intrinsic variable
    names, such as 'particle' or 'grid', inserts placeholder objects
    and propagates attribute access."""

    def __init__(self, grid, ptype):
        self.grid = grid
        self.ptype = ptype

    def visit_Name(self, node):
        if node.id == 'grid':
            return GridNode(self.grid, ccode='grid')
        elif node.id == 'particle':
            return ParticleNode(self.ptype, ccode='particle')
        else:
            return node

    def visit_Attribute(self, node):
        node.value = self.visit(node.value)
        if isinstance(node.value, IntrinsicNode):
            return getattr(node.value, node.attr)
        else:
            raise NotImplementedError("Cannot propagate attribute access to C-code")

    def visit_AugAssign(self, node):
        node.target = self.visit(node.target)
        node.op = self.visit(node.op)
        node.value = self.visit(node.value)

        # Capture p.lat/p.lon updates and insert p.xi/p.yi updates
        if isinstance(node.target, ParticleAttributeNode) \
           and node.target.ccode_index_var is not None:
            node = [node, node.target.pyast_index_update]
        return node

    def visit_Assign(self, node):
        node.targets = [self.visit(t) for t in node.targets]
        node.value = self.visit(node.value)

        # Capture p.lat/p.lon updates and insert p.xi/p.yi updates
        if isinstance(node.targets[0], ParticleAttributeNode) \
           and node.targets[0].ccode_index_var is not None:
            node = [node, node.targets[0].pyast_index_update]
        return node


class TupleSplitter(ast.NodeTransformer):
    """AST transformer that detects and splits Pythonic tuple
    assignments into multiple statements for conversion to C."""

    def visit_Assign(self, node):
        if isinstance(node.targets[0], ast.Tuple) \
           and isinstance(node.value, ast.Tuple):
            t_elts = node.targets[0].elts
            v_elts = node.value.elts
            if len(t_elts) != len(v_elts):
                raise AttributeError("Tuple lenghts in assignment do not agree")
            node = [ast.Assign() for _ in t_elts]
            for n, t, v in zip(node, t_elts, v_elts):
                n.targets = [t]
                n.value = v
        return node


class KernelGenerator(ast.NodeVisitor):
    """Code generator class that translates simple Python kernel
    functions into C functions by populating and accessing the `ccode`
    attriibute on nodes in the Python AST."""

    # Intrinsic variables that appear as function arguments
    kernel_vars = ['particle', 'grid', 'dt']

    def __init__(self, grid, ptype):
        self.grid = grid
        self.ptype = ptype

    def generate(self, pyfunc):
        # Parse the Python code into an AST
        self.py_ast = ast.parse(inspect.getsource(pyfunc.__code__))

        # Untangle Pythonic tuple-assignment statements
        self.py_ast = TupleSplitter().visit(self.py_ast)

        # Replace occurences of intrinsic objects in Python AST
        transformer = IntrinsicTransformer(self.grid, self.ptype)
        self.py_ast = transformer.visit(self.py_ast.body[0])

        # Generate C-code for all nodes in the Python AST
        self.visit(self.py_ast)
        self.ccode = self.py_ast.ccode

        # Derive local function variables and insert declaration
        funcvars = list(pyfunc.__code__.co_varnames)
        for kvar in self.kernel_vars:
            funcvars.remove(kvar)
        self.ccode.body.insert(0, c.Value("float", ", ".join(funcvars)))

        return self.ccode

    def visit_FunctionDef(self, node):
        # Create function declaration and argument list
        decl = c.Static(c.DeclSpecifier(c.Value("void", node.name), spec='inline'))
        U, V = (self.grid.U, self.grid.V)
        args = [c.Pointer(c.Value(self.ptype.name, "particle")), c.Value("float", "dt"),
                c.ArrayOf(c.Value("float", U.ccode_lon), count=U.lon.size),
                c.ArrayOf(c.Value("float", U.ccode_lat), count=U.lat.size),
{'fraction_non_alphanumeric': 0.11254253767622752, 'fraction_numerical': 0.0012153621779290229, 'mean_word_length': 2.439119531903037, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7204462', 'n_tokens_mistral': 3883, 'n_tokens_neox': 3721, 'n_words': 953}