Sergey B Kirpichev
61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
import ast
|
|
import io
|
|
import tokenize
|
|
|
|
|
|
class IntegerDivisionWrapper(ast.NodeTransformer):
|
|
"""Wrap all int divisions in a call to :class:`~fractions.Fraction`."""
|
|
|
|
def visit_BinOp(self, node):
|
|
def is_integer(x):
|
|
if isinstance(x, ast.Constant) and isinstance(x.value, int):
|
|
return True
|
|
if isinstance(x, ast.UnaryOp) and isinstance(x.op, (ast.USub,
|
|
ast.UAdd)):
|
|
return is_integer(x.operand)
|
|
if isinstance(x, ast.BinOp) and isinstance(x.op, (ast.Add,
|
|
ast.Sub,
|
|
ast.Mult,
|
|
ast.Pow)):
|
|
return is_integer(x.left) and is_integer(x.right)
|
|
return False
|
|
|
|
if isinstance(node.op, ast.Div) and all(map(is_integer,
|
|
[node.left, node.right])):
|
|
return ast.Call(ast.Name('Fraction', ast.Load()),
|
|
[node.left, node.right], [])
|
|
return self.generic_visit(node)
|
|
|
|
|
|
class _WrapFloats(ast.NodeTransformer):
|
|
"""Wrap float literals by calls to specified type."""
|
|
def __init__(self, lines, type):
|
|
super().__init__()
|
|
self.lines = lines
|
|
self.type = type
|
|
|
|
def visit_Constant(self, node):
|
|
if isinstance(node.value, (float, complex)):
|
|
line = self.lines[node.lineno - 1]
|
|
value = line[node.col_offset:node.end_col_offset]
|
|
is_complex = value.endswith(('j', 'J'))
|
|
if is_complex:
|
|
value = value[:-1]
|
|
value = ast.Constant(value)
|
|
value = ast.Call(ast.Name(self.type, ast.Load()), [value], [])
|
|
if is_complex:
|
|
value = ast.BinOp(left=value, op=ast.Mult(),
|
|
right=ast.Constant(1j))
|
|
return value
|
|
return node
|
|
|
|
|
|
def wrap_float_literals(lines):
|
|
"""Wraps all float/complex literals with mpmath classes."""
|
|
source = ''.join(lines)
|
|
tree = ast.parse(source)
|
|
tree = _WrapFloats(lines, 'mpf').visit(tree)
|
|
ast.fix_missing_locations(tree)
|
|
source = ast.unparse(tree)
|
|
return source.splitlines(keepends=True)
|