<commit_before>import itertools

# TODO: symbolic formatting: __add__ -> "+"

BINARY_LEVELS = {
        "__add__": 2,
        "__sub__": 2,

        "__mul__": 1,
        "__matmul__": 1,
        "__truediv__": 1,
        "__floordiv__": 1,
        "__mod__": 1,
        "__divmod__": 1,

        "__pow__": 0,

        "__lshift__": 3,
        "__rshift__": 3,

        "__and__": 4,
        "__xor__": 4,
        "__or__": 4,
        "__gt__": 5,
        "__lt__": 5,
        "__eq__": 5,
        "__ne__": 5,
        "__ge__": 5,
        "__le__": "5",

        "__getattr__": 0,
        "__getitem__": 0,
        }

BINARY_OPS = {
        "__add__": "+",
        "__sub__": "-",
        "__mul__": "*",
        "__matmul__": "@",
        "__truediv__": "/",
        "__floordiv__": "//",
        "__mod__": "%",
        "__divmod__": "IDK",
        "__pow__": "**",
        "__lshift__": "<<",
        "__rshift__": ">>",
        "__and__": "&",
        "__xor__": "^",
        "__or__": "|",
        "__gt__": ">",
        "__lt__": "<",
        "__eq__": "==",
        "__ne__": "!=",
        "__ge__": ">=",
        "__le__": "<=",
        "__getattr__": ".",
        "__getitem__": "[",
        }

UNARY_OPS = {
        "__invert__": "~",
        "__neg__": "-",
        "__pos__": "+",
        "__abs__": "ABS => "
        }

# TODO: can it just be put in binary ops? Special handling in Symbolic class?
ALL_OPS = {**BINARY_OPS, **UNARY_OPS}

#for k, v in BINARY_OPS.copy().items():
#    BINARY_OPS[k.replace("__", "__r", 1)] = v

for k, v in BINARY_LEVELS.copy().items():
    BINARY_LEVELS[k.replace("__", "__r", 1)] = v

class Formatter:
    def __init__(self): pass

    def format(self, call, pad = 0):
        """Return a Symbolic or Call back as a nice tree, with boxes for nodes."""

        fmt_block = "█─"
        fmt_pipe = "├─"
        
        # TODO: why are some nodes still symbolic?
        if isinstance(call, Symbolic):
            return self.format(strip_symbolic(call))

        if isinstance(call, MetaArg):
            return "_"

        if isinstance(call, Call):
            call_str = fmt_block + ALL_OPS.get(call.func, repr(call.func))

            args_str = [self.format(arg) for arg in call.args]

            # format keyword args, making sure "└─<key> = █─" aligns box's children
            kwargs_str = []
            for k, v in call.kwargs.items():
                kwargs_str.append(
                        k + " = " + self.format(v, pad = len(k) + 3)
                        )

            all_args = [*args_str, *kwargs_str]
            padded = []
            for ii, entry in enumerate(all_args):
                chunk = self.fmt_pipe(
                        entry,
                        final = ii == len(all_args) - 1,
                        pad = pad
                        )
                padded.append(chunk)

            return "".join([call_str, *padded])

        return repr(call)


    @staticmethod
    def fmt_pipe(x, final = False, pad = 0):
        if not final:
            connector = "│ " if not final else "  "
            prefix = "├─"
        else:
            connector = "  "
            prefix = "└─"

        connector = "\n" + " "*pad + connector
        prefix = "\n" + " "*pad + prefix
        # NOTE: because visiting is depth first, this is essentially prepending
        # the text to the left.
        return prefix + connector.join(x.splitlines())


# Calls
# =============================================================================

class Call:
    """
    Representation of python operations.

    This class is responsible for representing the pieces of a python expression,
    as a function, along with its args and kwargs.

    For example, "some_object.a" would be represented at the function "__getattr__",
    with the args `some_object`, and `"a"`.

    Args:
        func: name of the function called. Class methods are represented using the names
              they have when defined on the class.
        *args: arguments the func call uses.
        **kwargs: keyword arguments the func call uses.


    """
    def __init__(self, func, *args, **kwargs):
        self.func = func
        self.args = args
        self.kwargs = kwargs

    def __repr__(self):
        # TODO: format binary, unary, call, associative
        if self.func in BINARY_OPS:
            op_repr = BINARY_OPS[self.func]
            fmt = "({args[0]} {func} {args[1]})"
        elif self.func in UNARY_OPS:
            op_repr = UNARY_OPS[self.func]
            fmt = "({func}{args[0]})"
        elif self.func == "getattr":
            op_repr = "."
            fmt = "({args[0]}{func}{args[1]})"
        else:
            op_repr, *arg_str = map(repr, self.args)
            kwarg_str = (str(k) + " = " + repr(v) for k,v in self.kwargs.items())

            combined_arg_str = ",".join(itertools.chain(arg_str, kwarg_str))
            fmt = "{}({})".format(op_repr, combined_arg_str)
            return fmt

        return fmt.format(
                    func = op_repr or self.func,
                    args = self.args,
                    kwargs = self.kwargs
                    )

    def __call__(self, x):
        inst, *rest = (self.evaluate_calls(arg, x) for arg in self.args)
        kwargs = {k: self.evaluate_calls(v, x) for k, v in self.kwargs.items()}
        
        # TODO: temporary workaround, for when only __get_attribute__ is defined
        if self.func == "__getattr__":
            return getattr(inst, *rest)

        # in normal case, get method to call, and then call it
        return getattr(inst, self.func)(*rest, **kwargs)

    @staticmethod
    def evaluate_calls(arg, x):
        if isinstance(arg, Call): return arg(x)

        return arg

    def copy(self):
        args, kwargs = self.map_subcalls(lambda child: child.copy())
        return self.__class__(self.func, *args, **kwargs)

    def map_subcalls(self, f):
        new_args = tuple(f(arg) if isinstance(arg, Call) else arg for arg in self.args)
        new_kwargs = {k: f(v) if isinstance(v, Call) else v for k,v in self.kwargs.items()}

        return new_args, new_kwargs

    def iter_subcalls(self, f):
        yield from iter(arg for arg in self.args if instance(arg, Call))
        yield from iter(v for k,v in self.kwargs.items() if isinstance(v, Call))

    def op_vars(self, attr_calls = True):<commit_msg>feat: add map_replace method to Call<commit_after>import itertools

# TODO: symbolic formatting: __add__ -> "+"

BINARY_LEVELS = {
        "__add__": 2,
        "__sub__": 2,

        "__mul__": 1,
        "__matmul__": 1,
        "__truediv__": 1,
        "__floordiv__": 1,
        "__mod__": 1,
        "__divmod__": 1,

        "__pow__": 0,

        "__lshift__": 3,
        "__rshift__": 3,

        "__and__": 4,
        "__xor__": 4,
        "__or__": 4,
        "__gt__": 5,
        "__lt__": 5,
        "__eq__": 5,
        "__ne__": 5,
        "__ge__": 5,
        "__le__": "5",

        "__getattr__": 0,
        "__getitem__": 0,
        }

BINARY_OPS = {
        "__add__": "+",
        "__sub__": "-",
        "__mul__": "*",
        "__matmul__": "@",
        "__truediv__": "/",
        "__floordiv__": "//",
        "__mod__": "%",
        "__divmod__": "IDK",
        "__pow__": "**",
        "__lshift__": "<<",
        "__rshift__": ">>",
        "__and__": "&",
        "__xor__": "^",
        "__or__": "|",
        "__gt__": ">",
        "__lt__": "<",
        "__eq__": "==",
        "__ne__": "!=",
        "__ge__": ">=",
        "__le__": "<=",
        "__getattr__": ".",
        "__getitem__": "[",
        }

UNARY_OPS = {
        "__invert__": "~",
        "__neg__": "-",
        "__pos__": "+",
        "__abs__": "ABS => "
        }

# TODO: can it just be put in binary ops? Special handling in Symbolic class?
ALL_OPS = {**BINARY_OPS, **UNARY_OPS}

#for k, v in BINARY_OPS.copy().items():
#    BINARY_OPS[k.replace("__", "__r", 1)] = v

for k, v in BINARY_LEVELS.copy().items():
    BINARY_LEVELS[k.replace("__", "__r", 1)] = v

class Formatter:
    def __init__(self): pass

    def format(self, call, pad = 0):
        """Return a Symbolic or Call back as a nice tree, with boxes for nodes."""

        fmt_block = "█─"
        fmt_pipe = "├─"
        
        # TODO: why are some nodes still symbolic?
        if isinstance(call, Symbolic):
            return self.format(strip_symbolic(call))

        if isinstance(call, MetaArg):
            return "_"

        if isinstance(call, Call):
            call_str = fmt_block + ALL_OPS.get(call.func, repr(call.func))

            args_str = [self.format(arg) for arg in call.args]

            # format keyword args, making sure "└─<key> = █─" aligns box's children
            kwargs_str = []
            for k, v in call.kwargs.items():
                kwargs_str.append(
                        k + " = " + self.format(v, pad = len(k) + 3)
                        )

            all_args = [*args_str, *kwargs_str]
            padded = []
            for ii, entry in enumerate(all_args):
                chunk = self.fmt_pipe(
                        entry,
                        final = ii == len(all_args) - 1,
                        pad = pad
                        )
                padded.append(chunk)

            return "".join([call_str, *padded])

        return repr(call)


    @staticmethod
    def fmt_pipe(x, final = False, pad = 0):
        if not final:
            connector = "│ " if not final else "  "
            prefix = "├─"
        else:
            connector = "  "
            prefix = "└─"

        connector = "\n" + " "*pad + connector
        prefix = "\n" + " "*pad + prefix
        # NOTE: because visiting is depth first, this is essentially prepending
        # the text to the left.
        return prefix + connector.join(x.splitlines())


# Calls
# =============================================================================

class Call:
    """
    Representation of python operations.

    This class is responsible for representing the pieces of a python expression,
    as a function, along with its args and kwargs.

    For example, "some_object.a" would be represented at the function "__getattr__",
    with the args `some_object`, and `"a"`.

    Args:
        func: name of the function called. Class methods are represented using the names
              they have when defined on the class.
        *args: arguments the func call uses.
        **kwargs: keyword arguments the func call uses.


    """
    def __init__(self, func, *args, **kwargs):
        self.func = func
        self.args = args
        self.kwargs = kwargs

    def __repr__(self):
        # TODO: format binary, unary, call, associative
        if self.func in BINARY_OPS:
            op_repr = BINARY_OPS[self.func]
            fmt = "({args[0]} {func} {args[1]})"
        elif self.func in UNARY_OPS:
            op_repr = UNARY_OPS[self.func]
            fmt = "({func}{args[0]})"
        elif self.func == "getattr":
            op_repr = "."
            fmt = "({args[0]}{func}{args[1]})"
        else:
            op_repr, *arg_str = map(repr, self.args)
            kwarg_str = (str(k) + " = " + repr(v) for k,v in self.kwargs.items())

            combined_arg_str = ",".join(itertools.chain(arg_str, kwarg_str))
            fmt = "{}({})".format(op_repr, combined_arg_str)
            return fmt

        return fmt.format(
                    func = op_repr or self.func,
                    args = self.args,
                    kwargs = self.kwargs
                    )

    def __call__(self, x):
        inst, *rest = (self.evaluate_calls(arg, x) for arg in self.args)
        kwargs = {k: self.evaluate_calls(v, x) for k, v in self.kwargs.items()}
        
        # TODO: temporary workaround, for when only __get_attribute__ is defined
        if self.func == "__getattr__":
            return getattr(inst, *rest)

        # in normal case, get method to call, and then call it
        return getattr(inst, self.func)(*rest, **kwargs)

    @staticmethod
    def evaluate_calls(arg, x):
        if isinstance(arg, Call): return arg(x)

        return arg

    def copy(self):
        args, kwargs = self.map_subcalls(lambda child: child.copy())
        return self.__class__(self.func, *args, **kwargs)

    def map_subcalls(self, f):
        new_args = tuple(f(arg) if isinstance(arg, Call) else arg for arg in self.args)
        new_kwargs = {k: f(v) if isinstance(v, Call) else v for k,v in self.kwargs.items()}

        return new_args, new_kwargs

    def map_replace(self, f):
        args, kwargs = self.map_subcalls(f)
        return self.__class__(self.func, *args, **kwargs)

    def iter_subcalls(self, f):
        yield from iter(arg for arg in self.args if instance(arg, Call))
        yield from iter(v for k,v in self.kwargs.items() if isinstance(v, Call))

    def op_vars(self, attr_calls = True):
{'fraction_non_alphanumeric': 0.14908894379246448, 'fraction_numerical': 0.00509573810994441, 'mean_word_length': 1.8337344126011814, 'pattern_counts': {'":': 100, '<': 13, '<?xml version=': 0, '>': 17, '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': '3982266', 'n_tokens_mistral': 4183, 'n_tokens_neox': 4105, 'n_words': 1106}