wehub-resource-sync
91 lines
3.0 KiB
Python
91 lines
3.0 KiB
Python
import ast
|
|
import io
|
|
import re
|
|
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)
|
|
|
|
|
|
_HEXFLT_MATCHER = re.compile(r"""
|
|
(?: [^"' ]|^)[ ]*(?P<hexflt>
|
|
0x
|
|
[0-9a-z]+
|
|
(?: \.[0-9a-z]*)?
|
|
p(?:[+-])?[0-9]+
|
|
)
|
|
""", re.VERBOSE | re.IGNORECASE)
|
|
_BINFLT_MATCHER = re.compile(r"""
|
|
(?: [^"' ]|^)[ ]*(?P<binflt>
|
|
0b
|
|
[01]+
|
|
(?: \.[01]*)?
|
|
p(?:[+-])?[0-9]+
|
|
)
|
|
""", re.VERBOSE | re.IGNORECASE)
|
|
|
|
|
|
def wrap_hexbinfloats(lines):
|
|
new_lines = []
|
|
for line in lines:
|
|
for r in _HEXFLT_MATCHER.findall(line):
|
|
line = line.replace(r, 'mpf("' + r + '", base=16)')
|
|
for r in _BINFLT_MATCHER.findall(line):
|
|
line = line.replace(r, 'mpf("' + r + '", base=2)')
|
|
new_lines.append(line)
|
|
return new_lines
|