<commit_before>import ast
import sys
import traceback

from ..java import opcodes as JavaOpcodes
from .modules import Module
from .methods import MainFunction, InitMethod, to_java
from .structures import (
    IF, ELSE, END_IF,
    TRY, CATCH, FINALLY, END_TRY,
    START_LOOP, END_LOOP,
    ArgType,
    jump, OpcodePosition,
    AddToArgs, AddToKwargs
)
from .types.primitives import (
    ASTORE_name, ALOAD_name, free_name,
    ICONST_val, ISTORE_name, ILOAD_name
)
from .types import java, python
from .debug import (
    dump,
    # DEBUG, DEBUG_name
)


def is_call(node, name):
    return (
        # Node is a Call statement...
        isinstance(node, ast.Call)
        # ... where the function being invoked ...
        and isinstance(node.func, ast.Name)
        # ... is the provided name
        and (
            (isinstance(name, str) and node.func.id == name)
            or (isinstance(name, tuple) and node.func.id in name)
        )
    )


def node_visitor(fn):
    def dec(self, node, *args, **kwargs):
        try:
            if node.lineno != self._current_line:
                self._current_line = self.context.next_opcode_starts_line = node.lineno
        except AttributeError:
            pass
        self.context.next_resolve_list.append((node, OpcodePosition.START))
        fn(self, node, *args, **kwargs)
        self.context.next_resolve_list.append((node, OpcodePosition.END))
        self.context.next_resolve_list.append((node, OpcodePosition.NEXT))
    return dec


class NameVisitor(ast.NodeVisitor):
    def evaluate(self, root_node):
        self.names = [[]]
        if root_node:
            self.visit(root_node)
        return self

    @property
    def cls_name(self):
        return '$'.join('.'.join(n for n in group if n) for group in self.names)

    @property
    def ref_name(self):
        return '$'.join('/'.join(n for n in group if n) for group in self.names)

    @property
    def annotation(self):
        type_name = self.ref_name
        return type_name if type_name else 'org/python/Object'

    def visit_Name(self, node):
        self.names[-1].append(node.id)

    def visit_NameConstant(self, node):
        if node.value is None:
            self.names[-1].append("void")
        else:
            raise NotImplementedError("Unknown named constant %s" % node.value)

    def visit_Attribute(self, node):
        self.visit(node.value)
        self.names[-1].append(node.attr)

    def visit_Subscript(self, node):
        self.visit(node.value)
        self.names.append([])
        self.visit(node.slice)

    def visit_Index(self, node):
        self.visit(node.value)


class LocalsVisitor(ast.NodeVisitor):
    def __init__(self, context):
        self.context = context

    def visit_Name(self, node):
        if type(node.ctx) == ast.Store:
            if node.id not in self.context.local_vars:
                self.context.local_vars[node.id] = None

    def visit_Attribute(self, node):
        pass


class Visitor(ast.NodeVisitor):
    def __init__(self, namespace, filename, verbosity=1):
        super().__init__()
        self.namespace = namespace
        self.filename = filename
        self.verbosity = verbosity

        self._current_line = None
        self._context = []
        self._root_module = None

        self.current_exc_name = []

        self.symbol_namespace = {}
        self.code_objects = {}

    @property
    def context(self):
        return self._context[-1]

    @property
    def root_module(self):
        return self._root_module

    def push_context(self, block):
        self._context.append(block)
        block.visitor_setup()

    def pop_context(self):
        self.context.visitor_teardown()
        self._context.pop()

    def full_classref(self, name, default_prefix=None):
        return self.symbol_namespace.get(name, '.'.join([default_prefix, name])).replace('.', '/')

    def extract_code_objects(self, compiled):
        for obj in compiled.co_consts:
            if isinstance(obj, type(compiled)):
                if (obj.co_firstlineno, obj.co_name) in self.code_objects:
                    print(
                        'WARNING: multiple %s code objects found on line %s' % (
                            obj.co_name, obj.co_firstlineno
                        ),
                        file=sys.stderr
                    )
                self.code_objects[(obj.co_firstlineno, obj.co_name)] = obj
                self.extract_code_objects(obj)

    def visit(self, node):
        try:
            super().visit(node)
        except Exception as e:
            print(
                "Problem occurred in " + str(self.filename)
            )
            print('Node: ' + dump(node))
            raise
        return self._root_module

    def visit_Module(self, node):
        # Compile the module, and extract all the code objects
        compiled = compile(node, filename=self.filename, mode='exec')
        self.extract_code_objects(compiled)

        module = Module(self.namespace, self.filename)
        self.push_context(module)

        if self._root_module is None:
            self._root_module = module

        for child in node.body:
            self.visit(child)

        main = MainFunction(module)
        self.push_context(main)
        # No content, so pop right away.
        # We need to push to make sure setup/teardown
        # logic is invoked.
        self.pop_context()

        module.functions.append(main)

        self.pop_context()

    @node_visitor
    def visit_Interactive(self, node):
        # stmt* body):
        raise NotImplementedError('No handler for Interactive')

    @node_visitor
    def visit_Expr(self, node):
        self.generic_visit(node)

        # If the expression is a call, we need to ignore
        # any return value from the function.
        if isinstance(node.value, (ast.Call, ast.Attribute, ast.Str)):
            self.context.add_opcodes(
                JavaOpcodes.POP()
            )

    @node_visitor
    def visit_Suite(self, node):
        # stmt* body):
        raise NotImplementedError('No handler for Suite')

    # Statements
    @node_visitor
    def visit_FunctionDef(self, node):
        function = self._create_function(node, node.name, node.decorator_list)

        self.push_context(function)

        LocalsVisitor(function).visit(node)

        for child in node.body:
            self.visit(child)
                
        self.pop_context()
        
    @node_visitor
    def visit_ClassDef(self, node):
        # Construct a class.
        class_name = node.name

        name_visitor = NameVisitor()

        extends = None
        implements = []

        for keyword in node.keywords:
            key = keyword.arg
            value = keyword.value

            if key == "metaclass":
                raise Exception("Can't handle metaclasses")
            elif key == "extends":
                extends = name_visitor.evaluate(value).ref_name
            elif key == "implements":
                if isinstance(value, ast.List):
                    implements = [
                        name_visitor.visit(v).ref_name
                        for v in value.elts
                    ]
                else:
                    implements = [name_visitor.evaluate(value).ref_name]
            else:
                raise Exception("Unknown meta keyword " + str(key))

        self.context.add_opcodes(
            java.List(),
        )

        if extends:
            first_base = extends
        elif not node.bases:
            first_base = 'org/python/types/Object'
        else:
            first_base = None

        if first_base:
            self.context.add_opcodes(
                JavaOpcodes.DUP(),
                python.Type.for_class(first_base),
                java.List.add(),
            )

        for base in node.bases:
            self.context.add_opcodes(
                JavaOpcodes.DUP(),
            )
            self.visit(base)
            self.context.add_opcodes(
                java.List.add(),
            )

        klass = self.context.add_class(class_name, extends, implements)

        self.push_context(klass)
        for child in node.body:
            self.visit(child)
        self.pop_context()

        self.symbol_namespace[class_name] = klass.class_name

    @node_visitor
    def visit_Return(self, node):
        # expr? value):
        if node.value:
            self.visit(node.value)
        else:
            self.context.add_opcodes(python.NONE())
        self.context.add_opcodes(JavaOpcodes.ARETURN())
        # Record how deep we were when this return was added.
        self.context.opcodes[-1].needs_implicit_return = \
            self.context.has_nested_structure

    @node_visitor
    def visit_Delete(self, node):
        for target in node.targets:
            self.visit(target)
            if isinstance(target, ast.Attribute):
                self.context.add_opcodes(
                    python.Object.del_attr()
                )
            elif isinstance(target, ast.Subscript):
                self.context.add_opcodes(
                    python.Object.del_item()
                )
            elif isinstance(target, ast.Name):
                # delete is performed by visit(target) with context Del
                pass
            elif isinstance(target, (ast.Tuple, ast.List)):
                self.visit_Delete(ast.Delete(target.elts))
            else:
                raise NotImplementedError('No handler for Delete of type %s' % target)

    @node_visitor
    def visit_Assign(self, node):
        # Evaluate the value
        has_nested_yield = False
        for descendant_node in ast.walk(node.value):
            if isinstance(descendant_node, ast.Yield):
                has_nested_yield = True
                self.visit(descendant_node)
                break

        if has_nested_yield:
            self.context.load_name('<generator>')
            self.context.add_opcodes(
                JavaOpcodes.GETFIELD('org/python/types/Generator', 'message', 'Lorg/python/Object;')
            )

            # reset message to None after assignment
            self.context.load_name('<generator>')
            self.context.add_opcodes(
                JavaOpcodes.INVOKEVIRTUAL(
                    'org/python/types/Generator',
                    'reset_message',
                    args=[],
                    returns='V'
                )
            )
        else:
            self.visit(node.value)

        if len(node.targets) > 1:
            for target in node.targets:
                # Assign the value to the target
                self.context.add_opcodes(
                    JavaOpcodes.DUP()
                )
                self.visit(target)
            self.context.add_opcodes(
                JavaOpcodes.POP()
            )<commit_msg>Added support for BinOp, UOp, Call, etc on yield<commit_after>import ast
import sys
import traceback

from ..java import opcodes as JavaOpcodes
from .modules import Module
from .methods import MainFunction, InitMethod, to_java
from .structures import (
    IF, ELSE, END_IF,
    TRY, CATCH, FINALLY, END_TRY,
    START_LOOP, END_LOOP,
    ArgType,
    jump, OpcodePosition,
    AddToArgs, AddToKwargs
)
from .types.primitives import (
    ASTORE_name, ALOAD_name, free_name,
    ICONST_val, ISTORE_name, ILOAD_name
)
from .types import java, python
from .debug import (
    dump,
    # DEBUG, DEBUG_name
)


def is_call(node, name):
    return (
        # Node is a Call statement...
        isinstance(node, ast.Call)
        # ... where the function being invoked ...
        and isinstance(node.func, ast.Name)
        # ... is the provided name
        and (
            (isinstance(name, str) and node.func.id == name)
            or (isinstance(name, tuple) and node.func.id in name)
        )
    )


def node_visitor(fn):
    def dec(self, node, *args, **kwargs):
        try:
            if node.lineno != self._current_line:
                self._current_line = self.context.next_opcode_starts_line = node.lineno
        except AttributeError:
            pass
        self.context.next_resolve_list.append((node, OpcodePosition.START))
        fn(self, node, *args, **kwargs)
        self.context.next_resolve_list.append((node, OpcodePosition.END))
        self.context.next_resolve_list.append((node, OpcodePosition.NEXT))
    return dec


class NameVisitor(ast.NodeVisitor):
    def evaluate(self, root_node):
        self.names = [[]]
        if root_node:
            self.visit(root_node)
        return self

    @property
    def cls_name(self):
        return '$'.join('.'.join(n for n in group if n) for group in self.names)

    @property
    def ref_name(self):
        return '$'.join('/'.join(n for n in group if n) for group in self.names)

    @property
    def annotation(self):
        type_name = self.ref_name
        return type_name if type_name else 'org/python/Object'

    def visit_Name(self, node):
        self.names[-1].append(node.id)

    def visit_NameConstant(self, node):
        if node.value is None:
            self.names[-1].append("void")
        else:
            raise NotImplementedError("Unknown named constant %s" % node.value)

    def visit_Attribute(self, node):
        self.visit(node.value)
        self.names[-1].append(node.attr)

    def visit_Subscript(self, node):
        self.visit(node.value)
        self.names.append([])
        self.visit(node.slice)

    def visit_Index(self, node):
        self.visit(node.value)


class LocalsVisitor(ast.NodeVisitor):
    def __init__(self, context):
        self.context = context

    def visit_Name(self, node):
        if type(node.ctx) == ast.Store:
            if node.id not in self.context.local_vars:
                self.context.local_vars[node.id] = None

    def visit_Attribute(self, node):
        pass


class Visitor(ast.NodeVisitor):
    def __init__(self, namespace, filename, verbosity=1):
        super().__init__()
        self.namespace = namespace
        self.filename = filename
        self.verbosity = verbosity

        self._current_line = None
        self._context = []
        self._root_module = None

        self.current_exc_name = []

        self.symbol_namespace = {}
        self.code_objects = {}

    @property
    def context(self):
        return self._context[-1]

    @property
    def root_module(self):
        return self._root_module

    def push_context(self, block):
        self._context.append(block)
        block.visitor_setup()

    def pop_context(self):
        self.context.visitor_teardown()
        self._context.pop()

    def full_classref(self, name, default_prefix=None):
        return self.symbol_namespace.get(name, '.'.join([default_prefix, name])).replace('.', '/')

    def extract_code_objects(self, compiled):
        for obj in compiled.co_consts:
            if isinstance(obj, type(compiled)):
                if (obj.co_firstlineno, obj.co_name) in self.code_objects:
                    print(
                        'WARNING: multiple %s code objects found on line %s' % (
                            obj.co_name, obj.co_firstlineno
                        ),
                        file=sys.stderr
                    )
                self.code_objects[(obj.co_firstlineno, obj.co_name)] = obj
                self.extract_code_objects(obj)

    def visit(self, node):
        try:
            super().visit(node)
        except Exception as e:
            print(
                "Problem occurred in " + str(self.filename)
            )
            print('Node: ' + dump(node))
            raise
        return self._root_module

    def visit_Module(self, node):
        # Compile the module, and extract all the code objects
        compiled = compile(node, filename=self.filename, mode='exec')
        self.extract_code_objects(compiled)

        module = Module(self.namespace, self.filename)
        self.push_context(module)

        if self._root_module is None:
            self._root_module = module

        for child in node.body:
            self.visit(child)

        main = MainFunction(module)
        self.push_context(main)
        # No content, so pop right away.
        # We need to push to make sure setup/teardown
        # logic is invoked.
        self.pop_context()

        module.functions.append(main)

        self.pop_context()

    @node_visitor
    def visit_Interactive(self, node):
        # stmt* body):
        raise NotImplementedError('No handler for Interactive')

    @node_visitor
    def visit_Expr(self, node):
        self.generic_visit(node)

        # If the expression is a call, we need to ignore
        # any return value from the function.
        if isinstance(node.value, (ast.Call, ast.Attribute, ast.Str)):
            self.context.add_opcodes(
                JavaOpcodes.POP()
            )

    @node_visitor
    def visit_Suite(self, node):
        # stmt* body):
        raise NotImplementedError('No handler for Suite')

    # Statements
    @node_visitor
    def visit_FunctionDef(self, node):
        function = self._create_function(node, node.name, node.decorator_list)

        self.push_context(function)

        LocalsVisitor(function).visit(node)

        for child in node.body:
            self.visit(child)
                
        self.pop_context()
        
    @node_visitor
    def visit_ClassDef(self, node):
        # Construct a class.
        class_name = node.name

        name_visitor = NameVisitor()

        extends = None
        implements = []

        for keyword in node.keywords:
            key = keyword.arg
            value = keyword.value

            if key == "metaclass":
                raise Exception("Can't handle metaclasses")
            elif key == "extends":
                extends = name_visitor.evaluate(value).ref_name
            elif key == "implements":
                if isinstance(value, ast.List):
                    implements = [
                        name_visitor.visit(v).ref_name
                        for v in value.elts
                    ]
                else:
                    implements = [name_visitor.evaluate(value).ref_name]
            else:
                raise Exception("Unknown meta keyword " + str(key))

        self.context.add_opcodes(
            java.List(),
        )

        if extends:
            first_base = extends
        elif not node.bases:
            first_base = 'org/python/types/Object'
        else:
            first_base = None

        if first_base:
            self.context.add_opcodes(
                JavaOpcodes.DUP(),
                python.Type.for_class(first_base),
                java.List.add(),
            )

        for base in node.bases:
            self.context.add_opcodes(
                JavaOpcodes.DUP(),
            )
            self.visit(base)
            self.context.add_opcodes(
                java.List.add(),
            )

        klass = self.context.add_class(class_name, extends, implements)

        self.push_context(klass)
        for child in node.body:
            self.visit(child)
        self.pop_context()

        self.symbol_namespace[class_name] = klass.class_name

    @node_visitor
    def visit_Return(self, node):
        # expr? value):
        if node.value:
            self.visit(node.value)
        else:
            self.context.add_opcodes(python.NONE())
        self.context.add_opcodes(JavaOpcodes.ARETURN())
        # Record how deep we were when this return was added.
        self.context.opcodes[-1].needs_implicit_return = \
            self.context.has_nested_structure

    @node_visitor
    def visit_Delete(self, node):
        for target in node.targets:
            self.visit(target)
            if isinstance(target, ast.Attribute):
                self.context.add_opcodes(
                    python.Object.del_attr()
                )
            elif isinstance(target, ast.Subscript):
                self.context.add_opcodes(
                    python.Object.del_item()
                )
            elif isinstance(target, ast.Name):
                # delete is performed by visit(target) with context Del
                pass
            elif isinstance(target, (ast.Tuple, ast.List)):
                self.visit_Delete(ast.Delete(target.elts))
            else:
                raise NotImplementedError('No handler for Delete of type %s' % target)

    @node_visitor
    def visit_Assign(self, node):
        def push_message(node):
            # load message on stack
            self.context.load_name('<generator>')
            self.context.add_opcodes(
                JavaOpcodes.GETFIELD('org/python/types/Generator', 'message', 'Lorg/python/Object;')
            )

            # reset message to None after pushing it on stack
            self.context.load_name('<generator>')
            self.context.add_opcodes(
                JavaOpcodes.INVOKEVIRTUAL(
                    'org/python/types/Generator',
                    'reset_message',
                    args=[],
                    returns='V'
                )
            )

        # walk subtree with node.value as root to find any yield node
        nested_yield_node_found = False
        for descendant_node in ast.walk(node.value):
            if nested_yield_node_found:
                break

            if isinstance(descendant_node, ast.Yield):  # direct descendant, nothing much to process
                # visit the Yield node
                self.visit(descendant_node)
                push_message(descendant_node)
                nested_yield_node_found = True
            else:
                # go deeper into _fields of each descendant node
                for field_name, value in ast.iter_fields(descendant_node):
                    if isinstance(value, ast.Yield):
                        # visit the Yield node
                        nested_yield_node_found = True
                        self.visit(value)
                        push_message(value)

                        # convert the Yield node to Name node for expression evaluation
                        self.context.store_name('#msg-buffer-%x' % id(node))
                        new_node = ast.Name(
                            id='#msg-buffer-%x' % id(node),
                            ctx=ast.Load(),
                            lineno=value.lineno,
                            col_offset=value.col_offset
                        )
                        setattr(descendant_node, field_name, new_node)
                        break
                    elif isinstance(value, list):  # args or kwargs list
                        index = 0
                        for arg_node in value:
                            if isinstance(arg_node, ast.Yield):
                                # visit the Yield node
                                nested_yield_node_found = True
                                self.visit(arg_node)
                                push_message(arg_node)

                                # convert the Yield node to Name node for expression evaluation
                                # Use node.targets[0].id as temporary message storage
                                self.context.store_name('#msg-buffer-%x' % id(node))
                                new_node = ast.Name(
                                    id='#msg-buffer-%x' % id(node),
                                    ctx=ast.Load(),
                                    lineno=arg_node.lineno,
                                    col_offset=arg_node.col_offset
                                )
                                value[index] = new_node
                            index += 1
                        break

        if not isinstance(node.value, ast.Yield):
            # Evaluate the value
            self.visit(node.value)

        if len(node.targets) > 1:
            for target in node.targets:
                # Assign the value to the target
                self.context.add_opcodes(
                    JavaOpcodes.DUP()
                )
                self.visit(target)
            self.context.add_opcodes(
                JavaOpcodes.POP()
            )
{'fraction_non_alphanumeric': 0.09141844848185227, 'fraction_numerical': 0.0007003666625468628, 'mean_word_length': 1.6387650831612133, 'pattern_counts': {'":': 6, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3347187', 'n_tokens_mistral': 6822, 'n_tokens_neox': 6531, 'n_words': 1510}