"""Fortran/C symbolic expressions

References:

- J3/21-007: Draft Fortran 202x. https://j3-fortran.org/doc/year/21/21-007.pdf

Copyright 1999 -- 2011 Pearu Peterson all rights reserved.

Copyright 2011 -- present NumPy Developers.

Permission to use, modify, and distribute this software is given under the

terms of the NumPy License.

NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK.

"""

# To analyze Fortran expressions to solve dimensions specifications,

# for instances, we implement a minimal symbolic engine for parsing

# expressions into a tree of expression instances. As a first

# instance, we care only about arithmetic expressions involving

# integers and operations like addition (+), subtraction (-),

# multiplication (*), division (Fortran / is Python //, Fortran // is

# concatenate), and exponentiation (**). In addition, .pyf files may

# contain C expressions that support here is implemented as well.

#

# TODO: support logical constants (Op.BOOLEAN)

# TODO: support logical operators (.AND., ...)

# TODO: support defined operators (.MYOP., ...)

#

__all__ = ['Expr']

import re

import warnings

from enum import Enum

from math import gcd

class Language(Enum):

"""

Used as Expr.tostring language argument.

"""

Python = 0

Fortran = 1

C = 2

class Op(Enum):

"""

Used as Expr op attribute.

"""

INTEGER = 10

REAL = 12

COMPLEX = 15

STRING = 20

ARRAY = 30

SYMBOL = 40

TERNARY = 100

APPLY = 200

INDEXING = 210

CONCAT = 220

RELATIONAL = 300

TERMS = 1000

FACTORS = 2000

REF = 3000

DEREF = 3001

class RelOp(Enum):

"""

Used in Op.RELATIONAL expression to specify the function part.

"""

EQ = 1

NE = 2

LT = 3

LE = 4

GT = 5

GE = 6

@classmethod

def fromstring(cls, s, language=Language.C):

if language is Language.Fortran:

return {'.eq.': RelOp.EQ, '.ne.': RelOp.NE,

'.lt.': RelOp.LT, '.le.': RelOp.LE,

'.gt.': RelOp.GT, '.ge.': RelOp.GE}[s.lower()]

return {'==': RelOp.EQ, '!=': RelOp.NE, '<': RelOp.LT,

'<=': RelOp.LE, '>': RelOp.GT, '>=': RelOp.GE}[s]

def tostring(self, language=Language.C):

if language is Language.Fortran:

return {RelOp.EQ: '.eq.', RelOp.NE: '.ne.',

RelOp.LT: '.lt.', RelOp.LE: '.le.',

RelOp.GT: '.gt.', RelOp.GE: '.ge.'}[self]

return {RelOp.EQ: '==', RelOp.NE: '!=',

RelOp.LT: '<', RelOp.LE: '<=',

RelOp.GT: '>', RelOp.GE: '>='}[self]

class ArithOp(Enum):

"""

Used in Op.APPLY expression to specify the function part.

"""

POS = 1

NEG = 2

ADD = 3

SUB = 4

MUL = 5

DIV = 6

POW = 7

class OpError(Exception):

pass

class Precedence(Enum):

"""

Used as Expr.tostring precedence argument.

"""

ATOM = 0

POWER = 1

UNARY = 2

PRODUCT = 3

SUM = 4

LT = 6

EQ = 7

LAND = 11

LOR = 12

TERNARY = 13

ASSIGN = 14

TUPLE = 15

NONE = 100

integer_types = (int,)

number_types = (int, float)

def _pairs_add(d, k, v):

# Internal utility method for updating terms and factors data.

c = d.get(k)

if c is None:

d[k] = v

else:

c = c + v

if c:

d[k] = c

else:

del d[k]

class ExprWarning(UserWarning):

pass

def ewarn(message):

warnings.warn(message, ExprWarning, stacklevel=2)

class Expr:

"""Represents a Fortran expression as an op-data pair.

Expr instances are hashable and sortable.

"""

@staticmethod

def parse(s, language=Language.C):

"""Parse a Fortran expression to an Expr.

"""

return fromstring(s, language=language)

def __init__(self, op, data):

assert isinstance(op, Op)

# sanity checks

if op is Op.INTEGER:

# data is a 2-tuple of numeric object and a kind value

# (default is 4)

assert isinstance(data, tuple) and len(data) == 2

assert isinstance(data[0], int)

assert isinstance(data[1], (int, str)), data

elif op is Op.REAL:

# data is a 2-tuple of numeric object and a kind value

# (default is 4)

assert isinstance(data, tuple) and len(data) == 2

assert isinstance(data[0], float)

assert isinstance(data[1], (int, str)), data

elif op is Op.COMPLEX:

# data is a 2-tuple of constant expressions

assert isinstance(data, tuple) and len(data) == 2

elif op is Op.STRING:

# data is a 2-tuple of quoted string and a kind value

# (default is 1)

assert isinstance(data, tuple) and len(data) == 2

assert (isinstance(data[0], str)

and data[0][::len(data[0]) - 1] in ('""', "''", '@@'))

assert isinstance(data[1], (int, str)), data

elif op is Op.SYMBOL:

# data is any hashable object

assert hash(data) is not None

elif op in (Op.ARRAY, Op.CONCAT):

# data is a tuple of expressions

assert isinstance(data, tuple)

assert all(isinstance(item, Expr) for item in data), data

elif op in (Op.TERMS, Op.FACTORS):

# data is {<term|base>:<coeff|exponent>} where dict values

# are nonzero Python integers

assert isinstance(data, dict)

elif op is Op.APPLY:

# data is (<function>, <operands>, <kwoperands>) where

# operands are Expr instances

assert isinstance(data, tuple) and len(data) == 3

# function is any hashable object

assert hash(data[0]) is not None

assert isinstance(data[1], tuple)

assert isinstance(data[2], dict)

elif op is Op.INDEXING:

# data is (<object>, <indices>)

assert isinstance(data, tuple) and len(data) == 2

# function is any hashable object

assert hash(data[0]) is not None

elif op is Op.TERNARY:

# data is (<cond>, <expr1>, <expr2>)

assert isinstance(data, tuple) and len(data) == 3

elif op in (Op.REF, Op.DEREF):

# data is Expr instance

assert isinstance(data, Expr)

elif op is Op.RELATIONAL:

# data is (<relop>, <left>, <right>)

assert isinstance(data, tuple) and len(data) == 3

else:

raise NotImplementedError(

f'unknown op or missing sanity check: {op}')

self.op = op

self.data = data

def __eq__(self, other):

return (isinstance(other, Expr)

and self.op is other.op

and self.data == other.data)

def __hash__(self):

if self.op in (Op.TERMS, Op.FACTORS):

data = tuple(sorted(self.data.items()))

elif self.op is Op.APPLY:

data = self.data[:2] + tuple(sorted(self.data[2].items()))

else:

data = self.data

return hash((self.op, data))

def __lt__(self, other):

if isinstance(other, Expr):

if self.op is not other.op:

return self.op.value < other.op.value

if self.op in (Op.TERMS, Op.FACTORS):

return (tuple(sorted(self.data.items()))

< tuple(sorted(other.data.items())))

if self.op is Op.APPLY:

if self.data[:2] != other.data[:2]:

return self.data[:2] < other.data[:2]

return tuple(sorted(self.data[2].items())) < tuple(

sorted(other.data[2].items()))

return self.data < other.data

return NotImplemented

def __le__(self, other): return self == other or self < other

def __gt__(self, other): return not (self <= other)

def __ge__(self, other): return not (self < other)

def __repr__(self):

return f'{type(self).__name__}({self.op}, {self.data!r})'

def __str__(self):

return self.tostring()

def tostring(self, parent_precedence=Precedence.NONE,

language=Language.Fortran):

"""Return a string representation of Expr.

"""

if self.op in (Op.INTEGER, Op.REAL):

precedence = (Precedence.SUM if self.data[0] < 0

else Precedence.ATOM)

r = str(self.data[0]) + (f'_{self.data[1]}'

if self.data[1] != 4 else '')

elif self.op is Op.COMPLEX:

r = ', '.join(item.tostring(Precedence.TUPLE, language=language)

for item in self.data)

r = '(' + r + ')'

precedence = Precedence.ATOM

elif self.op is Op.SYMBOL:

precedence = Precedence.ATOM

r = str(self.data)

elif self.op is Op.STRING:

r = self.data[0]

if self.data[1] != 1:

r = self.data[1] + '_' + r

precedence = Precedence.ATOM

elif self.op is Op.ARRAY:

r = ', '.join(item.tostring(Precedence.TUPLE, language=language)

for item in self.data)

r = '[' + r + ']'

precedence = Precedence.ATOM

elif self.op is Op.TERMS:

terms = []

for term, coeff in sorted(self.data.items()):

if coeff < 0:

op = ' - '

coeff = -coeff

else:

op = ' + '

if coeff == 1:

term = term.tostring(Precedence.SUM, language=language)

elif term == as_number(1):

term = str(coeff)

else:

term = f'{coeff} * ' + term.tostring(

Precedence.PRODUCT, language=language)

if terms:

terms.append(op)

elif op == ' - ':

terms.append('-')

terms.append(term)

r = ''.join(terms) or '0'

precedence = Precedence.SUM if terms else Precedence.ATOM

elif self.op is Op.FACTORS:

factors = []

tail = []

for base, exp in sorted(self.data.items()):

op = ' * '

if exp == 1:

factor = base.tostring(Precedence.PRODUCT,

language=language)

elif language is Language.C:

if exp in range(2, 10):

factor = base.tostring(Precedence.PRODUCT,

language=language)

factor = ' * '.join([factor] * exp)

elif exp in range(-10, 0):

factor = base.tostring(Precedence.PRODUCT,

language=language)

tail += [factor] * -exp

continue

else:

factor = base.tostring(Precedence.TUPLE,

language=language)

factor = f'pow({factor}, {exp})'

else:

factor = base.tostring(Precedence.POWER,

language=language) + f' ** {exp}'

if factors:

factors.append(op)

factors.append(factor)

if tail:

if not factors:

factors += ['1']

factors += ['/', '(', ' * '.join(tail), ')']

r = ''.join(factors) or '1'

precedence = Precedence.PRODUCT if factors else Precedence.ATOM

elif self.op is Op.APPLY:

name, args, kwargs = self.data

if name is ArithOp.DIV and language is Language.C:

numer, denom = [arg.tostring(Precedence.PRODUCT,

language=language)

for arg in args]

r = f'{numer} / {denom}'

precedence = Precedence.PRODUCT

else:

args = [arg.tostring(Precedence.TUPLE, language=language)

for arg in args]

args += [k + '=' + v.tostring(Precedence.NONE)

for k, v in kwargs.items()]

r = f'{name}({", ".join(args)})'

precedence = Precedence.ATOM

elif self.op is Op.INDEXING:

name = self.data[0]

args = [arg.tostring(Precedence.TUPLE, language=language)

for arg in self.data[1:]]

r = f'{name}[{", ".join(args)}]'

precedence = Precedence.ATOM

elif self.op is Op.CONCAT:

args = [arg.tostring(Precedence.PRODUCT, language=language)

for arg in self.data]

r = " // ".join(args)

precedence = Precedence.PRODUCT

elif self.op is Op.TERNARY:

cond, expr1, expr2 = [a.tostring(Precedence.TUPLE,

language=language)

for a in self.data]

if language is Language.C:

r = f'({cond}?{expr1}:{expr2})'

elif language is Language.Python:

r = f'({expr1} if {cond} else {expr2})'

elif language is Language.Fortran:

r = f'merge({expr1}, {expr2}, {cond})'

else:

raise NotImplementedError(

f'tostring for {self.op} and {language}')

precedence = Precedence.ATOM

elif self.op is Op.REF:

r = '&' + self.data.tostring(Precedence.UNARY, language=language)

precedence = Precedence.UNARY

elif self.op is Op.DEREF:

r = '*' + self.data.tostring(Precedence.UNARY, language=language)

precedence = Precedence.UNARY

elif self.op is Op.RELATIONAL:

rop, left, right = self.data

precedence = (Precedence.EQ if rop in (RelOp.EQ, RelOp.NE)

else Precedence.LT)

left = left.tostring(precedence, language=language)

right = right.tostring(precedence, language=language)

rop = rop.tostring(language=language)

r = f'{left} {rop} {right}'

else:

raise NotImplementedError(f'tostring for op {self.op}')

if parent_precedence.value < precedence.value:

# If parent precedence is higher than operand precedence,

# operand will be enclosed in parenthesis.

return '(' + r + ')'

return r

def __pos__(self):

return self

def __neg__(self):

return self * -1

def __add__(self, other):

other = as_expr(other)

if isinstance(other, Expr):

if self.op is other.op:

if self.op in (Op.INTEGER, Op.REAL):

return as_number(

self.data[0] + other.data[0],

max(self.data[1], other.data[1]))

if self.op is Op.COMPLEX:

r1, i1 = self.data

r2, i2 = other.data

return as_complex(r1 + r2, i1 + i2)

if self.op is Op.TERMS:

r = Expr(self.op, dict(self.data))

for k, v in other.data.items():

_pairs_add(r.data, k, v)

return normalize(r)

if self.op is Op.COMPLEX and other.op in (Op.INTEGER, Op.REAL):

return self + as_complex(other)

elif self.op in (Op.INTEGER, Op.REAL) and other.op is Op.COMPLEX:

return as_complex(self) + other

elif self.op is Op.REAL and other.op is Op.INTEGER:

return self + as_real(other, kind=self.data[1])

elif self.op is Op.INTEGER and other.op is Op.REAL:

return as_real(self, kind=other.data[1]) + other

return as_terms(self) + as_terms(other)

return NotImplemented

def __radd__(self, other):

if isinstance(other, number_types):

return as_number(other) + self

return NotImplemented

def __sub__(self, other):

return self + (-other)

def __rsub__(self, other):

if isinstance(other, number_types):

return as_number(other) - self

return NotImplemented

def __mul__(self, other):

other = as_expr(other)

if isinstance(other, Expr):

if self.op is other.op:

if self.op in (Op.INTEGER, Op.REAL):

return as_number(self.data[0] * other.data[0],

max(self.data[1], other.data[1]))

elif self.op is Op.COMPLEX:

r1, i1 = self.data

r2, i2 = other.data

return as_complex(r1 * r2 - i1 * i2, r1 * i2 + r2 * i1)

if self.op is Op.FACTORS:

r = Expr(self.op, dict(self.data))

for k, v in other.data.items():

_pairs_add(r.data, k, v)

return normalize(r)

elif self.op is Op.TERMS:

r = Expr(self.op, {})

for t1, c1 in self.data.items():

for t2, c2 in other.data.items():

_pairs_add(r.data, t1 * t2, c1 * c2)

return normalize(r)

if self.op is Op.COMPLEX and other.op in (Op.INTEGER, Op.REAL):

return self * as_complex(other)

elif other.op is Op.COMPLEX and self.op in (Op.INTEGER, Op.REAL):

return as_complex(self) * other

elif self.op is Op.REAL and other.op is Op.INTEGER:

return self * as_real(other, kind=self.data[1])

elif self.op is Op.INTEGER and other.op is Op.REAL:

return as_real(self, kind=other.data[1]) * other

if self.op is Op.TERMS:

return self * as_terms(other)

elif other.op is Op.TERMS:

return as_terms(self) * other

return as_factors(self) * as_factors(other)

return NotImplemented

def __rmul__(self, other):

if isinstance(other, number_types):

return as_number(other) * self

return NotImplemented

def __pow__(self, other):

other = as_expr(other)

if isinstance(other, Expr):

if other.op is Op.INTEGER:

exponent = other.data[0]

# TODO: other kind not used

if exponent == 0:

return as_number(1)

if exponent == 1:

return self

if exponent > 0:

if self.op is Op.FACTORS:

r = Expr(self.op, {})

for k, v in self.data.items():

r.data[k] = v * exponent

return normalize(r)

return self * (self ** (exponent - 1))

elif exponent != -1:

return (self ** (-exponent)) ** -1

return Expr(Op.FACTORS, {self: exponent})

return as_apply(ArithOp.POW, self, other)

return NotImplemented

def __truediv__(self, other):

other = as_expr(other)

if isinstance(other, Expr):

# Fortran / is different from Python /:

# - `/` is a truncate operation for integer operands

return normalize(as_apply(ArithOp.DIV, self, other))

return NotImplemented

def __rtruediv__(self, other):

other = as_expr(other)

if isinstance(other, Expr):

return other / self

return NotImplemented

def __floordiv__(self, other):

other = as_expr(other)

if isinstance(other, Expr):

# Fortran // is different from Python //:

# - `//` is a concatenate operation for string operands

return normalize(Expr(Op.CONCAT, (self, other)))

return NotImplemented

def __rfloordiv__(self, other):

other = as_expr(other)

if isinstance(other, Expr):

return other // self

return NotImplemented

def __call__(self, *args, **kwargs):

# In Fortran, parenthesis () are use for both function call as

# well as indexing operations.

#

# TODO: implement a method for deciding when __call__ should

# return an INDEXING expression.

return as_apply(self, *map(as_expr, args),

**{k: as_expr(v) for k, v in kwargs.items()})

def __getitem__(self, index):

# Provided to support C indexing operations that .pyf files

# may contain.

index = as_expr(index)

if not isinstance(index, tuple):

index = index,

if len(index) > 1:

ewarn(f'C-index should be a single expression but got `{index}`')

return Expr(Op.INDEXING, (self,) + index)

def substitute(self, symbols_map):

"""Recursively substitute symbols with values in symbols map.

Symbols map is a dictionary of symbol-expression pairs.

"""

if self.op is Op.SYMBOL:

value = symbols_map.get(self)

if value is None:

return self

m = re.match(r'\A(@__f2py_PARENTHESIS_(\w+)_\d+@)\Z', self.data)

if m:

# complement to fromstring method

items, paren = m.groups()

if paren in ['ROUNDDIV', 'SQUARE']:

return as_array(value)

assert paren == 'ROUND', (paren, value)

return value

if self.op in (Op.INTEGER, Op.REAL, Op.STRING):

return self

if self.op in (Op.ARRAY, Op.COMPLEX):

return Expr(self.op, tuple(item.substitute(symbols_map)

for item in self.data))

if self.op is Op.CONCAT:

return normalize(Expr(self.op, tuple(item.substitute(symbols_map)

for item in self.data)))

if self.op is Op.TERMS:

r = None

for term, coeff in self.data.items():

if r is None:

r = term.substitute(symbols_map) * coeff

else:

r += term.substitute(symbols_map) * coeff

if r is None:

ewarn('substitute: empty TERMS expression interpreted as'

' int-literal 0')

return as_number(0)

return r

if self.op is Op.FACTORS:

r = None

for base, exponent in self.data.items():

if r is None:

r = base.substitute(symbols_map) ** exponent

else:

r *= base.substitute(symbols_map) ** exponent

if r is None:

ewarn('substitute: empty FACTORS expression interpreted'

' as int-literal 1')

return as_number(1)

return r

if self.op is Op.APPLY:

target, args, kwargs = self.data

if isinstance(target, Expr):

target = target.substitute(symbols_map)

args = tuple(a.substitute(symbols_map) for a in args)

kwargs = {k: v.substitute(symbols_map)

for k, v in kwargs.items()}

return normalize(Expr(self.op, (target, args, kwargs)))

if self.op is Op.INDEXING:

func = self.data[0]

if isinstance(func, Expr):

func = func.substitute(symbols_map)

args = tuple(a.substitute(symbols_map) for a in self.data[1:])

return normalize(Expr(self.op, (func,) + args))

if self.op is Op.TERNARY:

operands = tuple(a.substitute(symbols_map) for a in self.data)

return normalize(Expr(self.op, operands))

if self.op in (Op.REF, Op.DEREF):

return normalize(Expr(self.op, self.data.substitute(symbols_map)))

if self.op is Op.RELATIONAL:

rop, left, right = self.data

left = left.substitute(symbols_map)

right = right.substitute(symbols_map)

return normalize(Expr(self.op, (rop, left, right)))

raise NotImplementedError(f'substitute method for {self.op}: {self!r}')

def traverse(self, visit, *args, **kwargs):

"""Traverse expression tree with visit function.

The visit function is applied to an expression with given args

and kwargs.

Traverse call returns an expression returned by visit when not

None, otherwise return a new normalized expression with

traverse-visit sub-expressions.

"""

result = visit(self, *args, **kwargs)

if result is not None:

return result

if self.op in (Op.INTEGER, Op.REAL, Op.STRING, Op.SYMBOL):

return self

elif self.op in (Op.COMPLEX, Op.ARRAY, Op.CONCAT, Op.TERNARY):

return normalize(Expr(self.op, tuple(

item.traverse(visit, *args, **kwargs)

for item in self.data)))

elif self.op in (Op.TERMS, Op.FACTORS):

data = {}

for k, v in self.data.items():

k = k.traverse(visit, *args, **kwargs)

v = (v.traverse(visit, *args, **kwargs)

if isinstance(v, Expr) else v)

if k in data:

v = data[k] + v

data[k] = v

return normalize(Expr(self.op, data))

elif self.op is Op.APPLY:

obj = self.data[0]

func = (obj.traverse(visit, *args, **kwargs)

if isinstance(obj, Expr) else obj)

operands = tuple(operand.traverse(visit, *args, **kwargs)

for operand in self.data[1])

kwoperands = {k: v.traverse(visit, *args, **kwargs)

for k, v in self.data[2].items()}

return normalize(Expr(self.op, (func, operands, kwoperands)))

elif self.op is Op.INDEXING:

obj = self.data[0]

obj = (obj.traverse(visit, *args, **kwargs)

if isinstance(obj, Expr) else obj)

indices = tuple(index.traverse(visit, *args, **kwargs)

for index in self.data[1:])

return normalize(Expr(self.op, (obj,) + indices))

elif self.op in (Op.REF, Op.DEREF):

return normalize(Expr(self.op,

self.data.traverse(visit, *args, **kwargs)))

elif self.op is Op.RELATIONAL:

rop, left, right = self.data

left = left.traverse(visit, *args, **kwargs)

right = right.traverse(visit, *args, **kwargs)

return normalize(Expr(self.op, (rop, left, right)))

raise NotImplementedError(f'traverse method for {self.op}')

def contains(self, other):

"""Check if self contains other.

"""

found = []

def visit(expr, found=found):

if found:

return expr

elif expr == other:

found.append(1)

return expr

self.traverse(visit)

return len(found) != 0

def symbols(self):

"""Return a set of symbols contained in self.

"""

found = set()

def visit(expr, found=found):

if expr.op is Op.SYMBOL:

found.add(expr)

self.traverse(visit)

return found

def polynomial_atoms(self):

"""Return a set of expressions used as atoms in polynomial self.

"""

found = set()

def visit(expr, found=found):

if expr.op is Op.FACTORS:

for b in expr.data:

b.traverse(visit)

return expr

if expr.op in (Op.TERMS, Op.COMPLEX):

return

if expr.op is Op.APPLY and isinstance(expr.data[0], ArithOp):

if expr.data[0] is ArithOp.POW:

expr.data[1][0].traverse(visit)

return expr

return

if expr.op in (Op.INTEGER, Op.REAL):

return expr

found.add(expr)

if expr.op in (Op.INDEXING, Op.APPLY):

return expr

self.traverse(visit)

return found

def linear_solve(self, symbol):

"""Return a, b such that a * symbol + b == self.

If self is not linear with respect to symbol, raise RuntimeError.

"""

b = self.substitute({symbol: as_number(0)})

ax = self - b

a = ax.substitute({symbol: as_number(1)})

zero, _ = as_numer_denom(a * symbol - ax)

if zero != as_number(0):

raise RuntimeError(f'not a {symbol}-linear equation:'

f' {a} * {symbol} + {b} == {self}')

return a, b

def normalize(obj):

"""Normalize Expr and apply basic evaluation methods.

"""

if not isinstance(obj, Expr):

return obj

if obj.op is Op.TERMS:

d = {}

for t, c in obj.data.items():

if c == 0:

continue

if t.op is Op.COMPLEX and c != 1:

t = t * c

c = 1

if t.op is Op.TERMS:

for t1, c1 in t.data.items():

_pairs_add(d, t1, c1 * c)

else:

_pairs_add(d, t, c)

if len(d) == 0:

# TODO: determine correct kind

return as_number(0)

elif len(d) == 1:

(t, c), = d.items()

if c == 1:

return t

return Expr(Op.TERMS, d)

if obj.op is Op.FACTORS:

coeff = 1

d = {}

for b, e in obj.data.items():

if e == 0:

continue

if b.op is Op.TERMS and isinstance(e, integer_types) and e > 1:

# expand integer powers of sums

b = b * (b ** (e - 1))

e = 1

if b.op in (Op.INTEGER, Op.REAL):

if e == 1:

coeff *= b.data[0]

elif e > 0:

coeff *= b.data[0] ** e

else:

_pairs_add(d, b, e)

elif b.op is Op.FACTORS:

if e > 0 and isinstance(e, integer_types):

for b1, e1 in b.data.items():

_pairs_add(d, b1, e1 * e)

else:

_pairs_add(d, b, e)

else:

_pairs_add(d, b, e)

if len(d) == 0 or coeff == 0:

# TODO: determine correct kind

assert isinstance(coeff, number_types)

return as_number(coeff)

elif len(d) == 1:

(b, e), = d.items()

if e == 1:

t = b

else:

t = Expr(Op.FACTORS, d)

if coeff == 1:

return t

return Expr(Op.TERMS, {t: coeff})

elif coeff == 1:

return Expr(Op.FACTORS, d)

else:

return Expr(Op.TERMS, {Expr(Op.FACTORS, d): coeff})

if obj.op is Op.APPLY and obj.data[0] is ArithOp.DIV:

dividend, divisor = obj.data[1]

t1, c1 = as_term_coeff(dividend)

t2, c2 = as_term_coeff(divisor)

if isinstance(c1, integer_types) and isinstance(c2, integer_types):

g = gcd(c1, c2)

c1, c2 = c1 // g, c2 // g

else:

c1, c2 = c1 / c2, 1

if t1.op is Op.APPLY and t1.data[0] is ArithOp.DIV:

numer = t1.data[1][0] * c1

denom = t1.data[1][1] * t2 * c2

return as_apply(ArithOp.DIV, numer, denom)

if t2.op is Op.APPLY and t2.data[0] is ArithOp.DIV:

numer = t2.data[1][1] * t1 * c1

denom = t2.data[1][0] * c2

return as_apply(ArithOp.DIV, numer, denom)

d = dict(as_factors(t1).data)

for b, e in as_factors(t2).data.items():

_pairs_add(d, b, -e)

numer, denom = {}, {}

for b, e in d.items():

if e > 0:

numer[b] = e

else:

denom[b] = -e

numer = normalize(Expr(Op.FACTORS, numer)) * c1

denom = normalize(Expr(Op.FACTORS, denom)) * c2

if denom.op in (Op.INTEGER, Op.REAL) and denom.data[0] == 1:

# TODO: denom kind not used

return numer

return as_apply(ArithOp.DIV, numer, denom)

if obj.op is Op.CONCAT:

lst = [obj.data[0]]

for s in obj.data[1:]:

last = lst[-1]

if (

last.op is Op.STRING

and s.op is Op.STRING

and last.data[0][0] in '"\''

and s.data[0][0] == last.data[0][-1]

):

new_last = as_string(last.data[0][:-1] + s.data[0][1:],

max(last.data[1], s.data[1]))

lst[-1] = new_last

else:

lst.append(s)

if len(lst) == 1:

return lst[0]

return Expr(Op.CONCAT, tuple(lst))

if obj.op is Op.TERNARY:

cond, expr1, expr2 = map(normalize, obj.data)

if cond.op is Op.INTEGER:

return expr1 if cond.data[0] else expr2

return Expr(Op.TERNARY, (cond, expr1, expr2))

return obj

def as_expr(obj):

"""Convert non-Expr objects to Expr objects.

"""

if isinstance(obj, complex):

return as_complex(obj.real, obj.imag)

if isinstance(obj, number_types):

return as_number(obj)

if isinstance(obj, str):

# STRING expression holds string with boundary quotes, hence

# applying repr:

return as_string(repr(obj))

if isinstance(obj, tuple):

return tuple(map(as_expr, obj))

return obj

def as_symbol(obj):

"""Return object as SYMBOL expression (variable or unparsed expression).

"""

return Expr(Op.SYMBOL, obj)

def as_number(obj, kind=4):

"""Return object as INTEGER or REAL constant.

"""

if isinstance(obj, int):

return Expr(Op.INTEGER, (obj, kind))

if isinstance(obj, float):

return Expr(Op.REAL, (obj, kind))

if isinstance(obj, Expr):

if obj.op in (Op.INTEGER, Op.REAL):

return obj

raise OpError(f'cannot convert {obj} to INTEGER or REAL constant')

def as_integer(obj, kind=4):

"""Return object as INTEGER constant.

"""

if isinstance(obj, int):

return Expr(Op.INTEGER, (obj, kind))

if isinstance(obj, Expr):

if obj.op is Op.INTEGER:

return obj

raise OpError(f'cannot convert {obj} to INTEGER constant')

def as_real(obj, kind=4):

"""Return object as REAL constant.

"""

if isinstance(obj, int):

return Expr(Op.REAL, (float(obj), kind))

if isinstance(obj, float):

return Expr(Op.REAL, (obj, kind))

if isinstance(obj, Expr):

if obj.op is Op.REAL:

return obj

elif obj.op is Op.INTEGER:

return Expr(Op.REAL, (float(obj.data[0]), kind))

raise OpError(f'cannot convert {obj} to REAL constant')

def as_string(obj, kind=1):

"""Return object as STRING expression (string literal constant).

"""

return Expr(Op.STRING, (obj, kind))

def as_array(obj):

"""Return object as ARRAY expression (array constant).

"""

if isinstance(obj, Expr):

obj = obj,

return Expr(Op.ARRAY, obj)