<commit_before># This is very scratchy and supports only limited portion of Python functions.

import ast
import math
import inspect


class LatexifyVisitor(ast.NodeVisitor):
  def generic_visit(self, node):
    return str(node)

  def visit_Module(self, node):
    return self.visit(node.body[0])

  def visit_FunctionDef(self, node):
    name_str = r'\operatorname{' + str(node.name) + '}'
    arg_strs = [str(arg.arg) for arg in node.args.args]
    body_str = self.visit(node.body[0])
    return name_str + '(' + ', '.join(arg_strs) + r') \triangleq ' + body_str

  def visit_Return(self, node):
    return self.visit(node.value)

  def visit_Call(self, node):
    builtin_callees = {
        'abs': (r'\left|{', r'}\right|'),
        'math.acos': (r'\arccos{\left({', r'}\right)}'),
        'math.acosh': (r'\operatorname{arccosh}{\left({', r'}\right)}'),
        'math.asin': (r'\arcsin{\left({', r'}\right)}'),
        'math.asinh': (r'\operatorname{arcsinh}{\left({', r'}\right)}'),
        'math.atan': (r'\arctan{\left({', r'}\right)}'),
        'math.atanh': (r'\operatorname{arctanh}{\left({', r'}\right)}'),
        'math.ceil': (r'\left\lceil{', r'}\right\rceil'),
        'math.cos': (r'\cos{\left({', r'}\right)}'),
        'math.cosh': (r'\cosh{\left({', r'}\right)}'),
        'math.exp': (r'\exp{\left({', r'}\right)}'),
        'math.fabs': (r'\left|{', r'}\right|'),
        'math.factorial': (r'\left({', r'}\right)!'),
        'math.floor': (r'\left\lfloor{', r'}\right\rfloor'),
        'math.fsum': (r'\sum\left({', r'}\right)'),
        'math.gamma': (r'\Gamma\left({', r'}\right)'),
        'math.log': (r'\log{\left({', r'}\right)}'),
        'math.log10': (r'\log_{10}{\left({', r'}\right)}'),
        'math.log2': (r'\log_{2}{\left({', r'}\right)}'),
        'math.prod': (r'\prod \left({', r'}\right)'),
        'math.sin': (r'\sin{\left({', r'}\right)}'),
        'math.sinh': (r'\sinh{\left({', r'}\right)}'),
        'math.sqrt': (r'\sqrt{', '}'),
        'math.tan': (r'\tan{\left({', r'}\right)}'),
        'math.tanh': (r'\tanh{\left({', r'}\right)}'),
        'sum': (r'\sum \left({', r'}\right)'),
    }

    callee_str = self.visit(node.func)
    if callee_str in builtin_callees:
      lstr, rstr = builtin_callees[callee_str]
    else:
      if callee_str.startswith('math.'):
        callee_str = callee_str[5:]
      lstr = r'\operatorname{' + callee_str + '}\left('
      rstr = r'\right)'

    arg_strs = [self.visit(arg) for arg in node.args]
    return lstr + ', '.join(arg_strs) + rstr

  def visit_Attribute(self, node):
    vstr = self.visit(node.value)
    astr = str(node.attr)
    return vstr + '.' + astr

  def visit_Name(self, node):
    return str(node.id)

  def visit_Num(self, node):
    return str(node.n)

  def visit_UnaryOp(self, node):
    def _wrap(child):
      latex = self.visit(child)
      if isinstance(child, ast.BinOp) and isinstance(child.op, (ast.Add, ast.Sub)):
        return '(' + latex + ')'
      return latex

    reprs = {
        ast.UAdd: (lambda: _wrap(node.operand)),
        ast.USub: (lambda: '-' + _wrap(node.operand)),
    }

    if type(node.op) in reprs:
      return reprs[type(node.op)]()
    else:
      return r'\operatorname{unknown\_uniop}(' + vstr + ')'

  def visit_BinOp(self, node):
    priority = {
        ast.Add: 10,
        ast.Sub: 10,
        ast.Mult: 20,
        ast.Div: 20,
        ast.Pow: 30,
    }

    def _unwrap(child):
      return self.visit(child)

    def _wrap(child):
      latex = _unwrap(child)
      if isinstance(child, ast.BinOp):
        cp = priority[type(child.op)] if type(child.op) in priority else 100
        pp = priority[type(node.op)] if type(node.op) in priority else 100
        if cp < pp:
          return '(' + latex + ')'
      return latex

    l = node.left
    r = node.right
    reprs = {
        ast.Add: (lambda: _wrap(l) + ' + ' + _wrap(r)),
        ast.Sub: (lambda: _wrap(l) + ' - ' + _wrap(r)),
        ast.Mult: (lambda: _wrap(l) + _wrap(r)),
        ast.Div: (lambda: r'\frac{' + _unwrap(l) + '}{' + _unwrap(r) + '}'),
        ast.Pow: (lambda: _wrap(l) + '^{' + _unwrap(r) + '}'),
    }

    if type(node.op) in reprs:
      return reprs[type(node.op)]()
    else:
      return r'\operatorname{unknown\_binop}(' + lstr + ', ' + rstr + ')'

  def visit_Compare(self, node):
    lstr = self.visit(node.left)
    rstr = self.visit(node.comparators[0])

    if isinstance(node.ops[0], ast.Eq):
      return lstr + '=' + rstr
    else:
      return r'\operatorname{unknown\_comparator}(' + lstr + ', ' + rstr + ')'

  def visit_If(self, node):
    cond_str = self.visit(node.test)
    tstr = self.visit(node.body[0])
    fstr = self.visit(node.orelse[0])
    return r'\left\{ \begin{array}{ll} ' + tstr + r', & \mathrm{if} \ ' + cond_str + r' \\ ' + fstr + r', & \mathrm{otherwise} \end{array} \right.'


def get_latex(fn):
  return LatexifyVisitor().visit(ast.parse(inspect.getsource(fn)))


def with_latex(fn):

  class _LatexifiedFunction:
    def __init__(self, fn):<commit_msg>Fixes bug on unknown BinOp.<commit_after># This is very scratchy and supports only limited portion of Python functions.

import ast
import math
import inspect


class LatexifyVisitor(ast.NodeVisitor):
  def generic_visit(self, node):
    return str(node)

  def visit_Module(self, node):
    return self.visit(node.body[0])

  def visit_FunctionDef(self, node):
    name_str = r'\operatorname{' + str(node.name) + '}'
    arg_strs = [str(arg.arg) for arg in node.args.args]
    body_str = self.visit(node.body[0])
    return name_str + '(' + ', '.join(arg_strs) + r') \triangleq ' + body_str

  def visit_Return(self, node):
    return self.visit(node.value)

  def visit_Call(self, node):
    builtin_callees = {
        'abs': (r'\left|{', r'}\right|'),
        'math.acos': (r'\arccos{\left({', r'}\right)}'),
        'math.acosh': (r'\operatorname{arccosh}{\left({', r'}\right)}'),
        'math.asin': (r'\arcsin{\left({', r'}\right)}'),
        'math.asinh': (r'\operatorname{arcsinh}{\left({', r'}\right)}'),
        'math.atan': (r'\arctan{\left({', r'}\right)}'),
        'math.atanh': (r'\operatorname{arctanh}{\left({', r'}\right)}'),
        'math.ceil': (r'\left\lceil{', r'}\right\rceil'),
        'math.cos': (r'\cos{\left({', r'}\right)}'),
        'math.cosh': (r'\cosh{\left({', r'}\right)}'),
        'math.exp': (r'\exp{\left({', r'}\right)}'),
        'math.fabs': (r'\left|{', r'}\right|'),
        'math.factorial': (r'\left({', r'}\right)!'),
        'math.floor': (r'\left\lfloor{', r'}\right\rfloor'),
        'math.fsum': (r'\sum\left({', r'}\right)'),
        'math.gamma': (r'\Gamma\left({', r'}\right)'),
        'math.log': (r'\log{\left({', r'}\right)}'),
        'math.log10': (r'\log_{10}{\left({', r'}\right)}'),
        'math.log2': (r'\log_{2}{\left({', r'}\right)}'),
        'math.prod': (r'\prod \left({', r'}\right)'),
        'math.sin': (r'\sin{\left({', r'}\right)}'),
        'math.sinh': (r'\sinh{\left({', r'}\right)}'),
        'math.sqrt': (r'\sqrt{', '}'),
        'math.tan': (r'\tan{\left({', r'}\right)}'),
        'math.tanh': (r'\tanh{\left({', r'}\right)}'),
        'sum': (r'\sum \left({', r'}\right)'),
    }

    callee_str = self.visit(node.func)
    if callee_str in builtin_callees:
      lstr, rstr = builtin_callees[callee_str]
    else:
      if callee_str.startswith('math.'):
        callee_str = callee_str[5:]
      lstr = r'\operatorname{' + callee_str + '}\left('
      rstr = r'\right)'

    arg_strs = [self.visit(arg) for arg in node.args]
    return lstr + ', '.join(arg_strs) + rstr

  def visit_Attribute(self, node):
    vstr = self.visit(node.value)
    astr = str(node.attr)
    return vstr + '.' + astr

  def visit_Name(self, node):
    return str(node.id)

  def visit_Num(self, node):
    return str(node.n)

  def visit_UnaryOp(self, node):
    def _wrap(child):
      latex = self.visit(child)
      if isinstance(child, ast.BinOp) and isinstance(child.op, (ast.Add, ast.Sub)):
        return '(' + latex + ')'
      return latex

    reprs = {
        ast.UAdd: (lambda: _wrap(node.operand)),
        ast.USub: (lambda: '-' + _wrap(node.operand)),
    }

    if type(node.op) in reprs:
      return reprs[type(node.op)]()
    else:
      return r'\operatorname{unknown\_uniop}(' + vstr + ')'

  def visit_BinOp(self, node):
    priority = {
        ast.Add: 10,
        ast.Sub: 10,
        ast.Mult: 20,
        ast.Div: 20,
        ast.Pow: 30,
    }

    def _unwrap(child):
      return self.visit(child)

    def _wrap(child):
      latex = _unwrap(child)
      if isinstance(child, ast.BinOp):
        cp = priority[type(child.op)] if type(child.op) in priority else 100
        pp = priority[type(node.op)] if type(node.op) in priority else 100
        if cp < pp:
          return '(' + latex + ')'
      return latex

    l = node.left
    r = node.right
    reprs = {
        ast.Add: (lambda: _wrap(l) + ' + ' + _wrap(r)),
        ast.Sub: (lambda: _wrap(l) + ' - ' + _wrap(r)),
        ast.Mult: (lambda: _wrap(l) + _wrap(r)),
        ast.Div: (lambda: r'\frac{' + _unwrap(l) + '}{' + _unwrap(r) + '}'),
        ast.Pow: (lambda: _wrap(l) + '^{' + _unwrap(r) + '}'),
    }

    if type(node.op) in reprs:
      return reprs[type(node.op)]()
    else:
      return r'\operatorname{unknown\_binop}(' + _unwrap(l) + ', ' + _unwrap(r) + ')'

  def visit_Compare(self, node):
    lstr = self.visit(node.left)
    rstr = self.visit(node.comparators[0])

    if isinstance(node.ops[0], ast.Eq):
      return lstr + '=' + rstr
    else:
      return r'\operatorname{unknown\_comparator}(' + lstr + ', ' + rstr + ')'

  def visit_If(self, node):
    cond_str = self.visit(node.test)
    tstr = self.visit(node.body[0])
    fstr = self.visit(node.orelse[0])
    return r'\left\{ \begin{array}{ll} ' + tstr + r', & \mathrm{if} \ ' + cond_str + r' \\ ' + fstr + r', & \mathrm{otherwise} \end{array} \right.'


def get_latex(fn):
  return LatexifyVisitor().visit(ast.parse(inspect.getsource(fn)))


def with_latex(fn):

  class _LatexifiedFunction:
    def __init__(self, fn):
{'fraction_non_alphanumeric': 0.2191089108910891, 'fraction_numerical': 0.0057425742574257425, 'mean_word_length': 2.7108743570903746, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5139207', 'n_tokens_mistral': 3875, 'n_tokens_neox': 3706, 'n_words': 723}