pine2pyne 0.2.4__py3-none-any.whl

pine2pyne/emit/pynecomp/passes/rewrites.py

@@ -0,0 +1,215 @@
+"""The PyneComp rewrites, one class per difference.
+
+Each docstring says what PyneComp does and how often ``assets/pyne`` shows it.
+"""
+from __future__ import annotations
+
+import ast
+import re
+
+from .base import PassContext, function_scopes, main_function, own_nodes, statements_in_order
+
+_HEX_COLOR = re.compile(r'#(?:[0-9a-fA-F]{6}|[0-9a-fA-F]{8})')
+
+
+def _double_quoted(text: str) -> str:
+    return '"' + text.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n') + '"'
+
+
+class DeclarationCallShapePass:
+    """``@script.indicator(...)`` keeps the call shape of the Pine declaration:
+    the arguments the source passed positionally stay positional
+    (``indicator('T', 'S', overlay=true)`` -> ``('T', 'S', overlay=True)``), a
+    title passed by keyword stays ``title=``, and a Pine v6 title keeps its quotes."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        main = main_function(tree)
+        declaration = getattr(context.unit.pine_script, 'script_decl', None)
+        if main is None or declaration is None:
+            return
+        for decorator in main.decorator_list:
+            if (isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Attribute)
+                    and isinstance(decorator.func.value, ast.Name) and decorator.func.value.id == 'script'):
+                self._reshape(decorator, declaration, context)
+
+    @staticmethod
+    def _reshape(call: ast.Call, declaration, context: PassContext) -> None:
+        positional = getattr(declaration, 'positional_count', 0)
+        order = tuple(getattr(declaration, 'keyword_order', ()))
+        title = call.args[0] if call.args else None
+        if title is None:  # a title that is an expression: the first keyword
+            if positional:
+                call.args = [keyword.value for keyword in call.keywords[:positional]]
+                call.keywords = call.keywords[positional:]
+            return
+        if positional == 0:
+            call.args = []
+            at = order.index('title') if 'title' in order else 0
+            call.keywords = call.keywords[:at] + [ast.keyword(arg='title', value=title)] + call.keywords[at:]
+        elif positional > 1:
+            taken = call.keywords[:positional - 1]
+            call.args = [title] + [keyword.value for keyword in taken]
+            call.keywords = call.keywords[positional - 1:]
+        if (title is not None and getattr(declaration, 'title_double_quoted', False)
+                and context.pine_version >= 6 and isinstance(title, ast.Constant)):
+            context.spellings[id(title)] = _double_quoted(title.value)
+
+
+class StrategyDefaultsPass:
+    """A Pine v4/v5 ``strategy()`` is compiled with its version's defaults made
+    explicit, appended after the source's own arguments in this order (every
+    v4/v5 strategy of the reference): ``margin_long=0, margin_short=0,
+    initial_capital=100000`` (v4) or ``1000000`` (v5), ``default_qty_type=strategy.fixed,
+    default_qty_value=1, close_by_id=True``. An argument the source gives keeps its value."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        if context.pine_version >= 6:
+            return
+        main = main_function(tree)
+        for decorator in main.decorator_list if main else ():
+            if (isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Attribute)
+                    and decorator.func.attr == 'strategy'):
+                self._add_defaults(decorator, context.pine_version)
+
+    @staticmethod
+    def _add_defaults(call: ast.Call, version: int) -> None:
+        given = {keyword.arg for keyword in call.keywords}
+        defaults = [
+            ('margin_long', ast.Constant(0)),
+            ('margin_short', ast.Constant(0)),
+            ('initial_capital', ast.Constant(100000 if version <= 4 else 1000000)),
+            ('default_qty_type', ast.Attribute(value=ast.Name(id='strategy', ctx=ast.Load()), attr='fixed',
+                                               ctx=ast.Load())),
+            ('default_qty_value', ast.Constant(1)),
+            ('close_by_id', ast.Constant(True)),
+        ]
+        added = [ast.keyword(arg=name, value=value) for name, value in defaults if name not in given]
+        # PyneComp reconciles the imports before it adds these: `strategy.fixed` alone imports nothing
+        for keyword in added:
+            for node in ast.walk(keyword.value):
+                node.added_default = True
+        tail = [keyword for keyword in call.keywords if keyword.arg == 'na_bool']
+        call.keywords = [keyword for keyword in call.keywords if keyword.arg != 'na_bool'] + added + tail
+
+
+class NaBoolFlagPass:
+    """``na_bool=True`` (Pine's three-state bool) is kept for a Pine v4 script
+    (all 295 of the reference have it). A Pine v5 script keeps it only when it
+    reads a named bool where it may be na: ``b[1]`` or ``na(b)`` (probe scripts
+    compiled by PyneComp v6.0.80; 82 of the 327 v5 scripts of the reference)."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        if context.pine_version >= 6 and getattr(context.unit.transformed, 'reads_na_of_bool', False):
+            main = main_function(tree)  # Pine v6: only `na()` of a bool collection marks it (probe nb6_02)
+            for decorator in main.decorator_list if main else ():
+                if isinstance(decorator, ast.Call) and not any(k.arg == 'na_bool' for k in decorator.keywords):
+                    decorator.keywords.append(ast.keyword(arg='na_bool', value=ast.Constant(True)))
+            return
+        if context.pine_version != 5 or getattr(context.unit.transformed, 'reads_na_of_bool', False):
+            return
+        main = main_function(tree)
+        for decorator in main.decorator_list if main else ():
+            if isinstance(decorator, ast.Call):
+                decorator.keywords = [keyword for keyword in decorator.keywords if keyword.arg != 'na_bool']
+
+
+class MathNamespacePass:
+    """``abs``, ``max``, ``min``, ``pow`` and ``round`` are always called through
+    ``math.`` (the reference never calls the Python builtins: 956 call sites)."""
+
+    NAMES = frozenset({'abs', 'max', 'min', 'pow', 'round'})
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        bound = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store)}
+        bound |= {node.name for node in ast.walk(tree) if isinstance(node, (ast.FunctionDef, ast.ClassDef))}
+        for node in ast.walk(tree):
+            if (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)
+                    and node.func.id in self.NAMES and node.func.id not in bound):
+                node.func = ast.copy_location(ast.Attribute(
+                    value=ast.Name(id='math', ctx=ast.Load()), attr=node.func.id, ctx=ast.Load()), node.func)
+
+
+class FlatBoolOpPass:
+    """``a and (b and c)`` is one ``and`` of three operands: the reference has no
+    ``and`` directly inside an ``and`` (nor ``or`` in ``or``), though its sources
+    write them (119 ``and (... and ...)`` groups)."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        for node in ast.walk(tree):
+            if isinstance(node, ast.BoolOp):
+                node.values = self._flat(node)
+
+    def _flat(self, node: ast.BoolOp) -> list:
+        values = []
+        for value in node.values:
+            if isinstance(value, ast.BoolOp) and type(value.op) is type(node.op):
+                values.extend(self._flat(value))
+            else:
+                values.append(value)
+        return values
+
+
+class ArithmeticReassociationPass:
+    """Pine v4/v5: a parenthesised right operand is re-associated to the left where
+    the operators allow it: ``a + (b - c)`` is ``a + b - c`` and ``a * (b * c)`` is
+    ``a * b * c``; ``a * (b / c)`` and every right operand of ``-`` or ``/`` keep
+    their parentheses. Pine v6 keeps the source's grouping. Probe scripts compiled
+    by PyneComp v6.0.80 (``tmp/probes/par4..6``)."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        if context.pine_version >= 6:
+            return
+        _Reassociation().visit(tree)
+
+
+class _Reassociation(ast.NodeTransformer):
+    def visit_BinOp(self, node: ast.BinOp) -> ast.AST:
+        self.generic_visit(node)
+        return _reassociated(node)
+
+
+def _reassociated(node: ast.BinOp) -> ast.BinOp:
+    right = node.right
+    joins = (isinstance(node.op, ast.Add) and isinstance(right, ast.BinOp) and isinstance(right.op, (ast.Add, ast.Sub))
+             or isinstance(node.op, ast.Mult) and isinstance(right, ast.BinOp) and isinstance(right.op, ast.Mult))
+    if not joins:
+        return node
+    left = _reassociated(ast.copy_location(ast.BinOp(left=node.left, op=node.op, right=right.left), node))
+    return _reassociated(ast.copy_location(ast.BinOp(left=left, op=right.op, right=right.right), node))
+
+
+class HexColorPass:
+    """A Pine colour literal is ``color.new('#RRGGBB')``, upper case
+    (all 5410 hex colours of the reference)."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        tree.body[:] = [self._visit(node) for node in tree.body]
+
+    def _visit(self, node):
+        if isinstance(node, (ast.Import, ast.ImportFrom)):
+            return node
+        return _HexColorRewriter().visit(node)
+
+
+class _HexColorRewriter(ast.NodeTransformer):
+    def visit_Constant(self, node: ast.Constant):
+        if isinstance(node.value, str) and _HEX_COLOR.fullmatch(node.value):
+            spelled = node.value.upper()
+            if len(spelled) == 9 and spelled.endswith('FF'):
+                spelled = spelled[:7]  # an opaque alpha is left out: '#00FF00FF' is '#00FF00'
+            call = ast.Call(
+                func=ast.Attribute(value=ast.Name(id='color', ctx=ast.Load()), attr='new', ctx=ast.Load()),
+                args=[ast.Constant(value=spelled)], keywords=[])
+            return ast.copy_location(call, node)
+        return node
+
+
+__all__ = [
+    'DeclarationCallShapePass',
+    'FlatBoolOpPass',
+    'MathNamespacePass',
+    'NaBoolFlagPass',
+    'StrategyDefaultsPass',
+    'HexColorPass',
+    'statements_in_order',
+]

pine2pyne/emit/pynecomp/printer.py

@@ -0,0 +1,288 @@
+"""Module layout of PyneComp output, printed from a Python AST.
+
+The printer owns layout only: what a statement *is* comes from the code
+generator; how it is spread over lines, indented and separated comes from
+here. Expressions go through :class:`~.unparser.ExpressionPrinter`, the
+import block through :class:`~.imports.ImportBlockFormatter`, blank lines in
+function bodies through a :class:`~.blank_lines.BlankLinePolicy`.
+"""
+from __future__ import annotations
+
+import ast
+import re
+from dataclasses import dataclass, field
+from typing import Callable, Mapping
+
+from .banner import PyneCompBanner
+from .blank_lines import AfterDefinitionPolicy, BlankLinePolicy, SourceLayoutPolicy
+from .imports import ImportBlockFormatter
+from .unparser import ExpressionPrinter, header_strings, source_quoted_strings
+
+INDENT = '    '
+_DEFINITIONS = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
+
+
+@dataclass(frozen=True)
+class LayoutSource:
+    """What the printer reads: the generated code and where it came from."""
+
+    python_source: str
+    pine_source: str = ''
+    pine_version: int | None = None
+    #: Python line number -> Pine line number of the statement printed there
+    origins: Mapping[int, int] = field(default_factory=dict)
+
+    @property
+    def keeps_source_layout(self) -> bool:
+        """Pine v6 keeps its blank lines and string quotes; earlier versions do not."""
+        return (self.pine_version or 6) >= 6
+
+
+class _ModuleWriter:
+    """Prints one module; created per render, so the printer stays stateless."""
+
+    def __init__(self, expression: Callable[[ast.AST], str], blank_lines: BlankLinePolicy,
+                 written_as_elif: Callable[[ast.If], bool] = lambda node: True,
+                 blank_in_source: Callable[[ast.stmt], bool] = lambda node: False) -> None:
+        self._expression = expression
+        self._blank_lines = blank_lines
+        self._written_as_elif = written_as_elif
+        self._blank_in_source = blank_in_source
+
+    # -- functions ---------------------------------------------------------
+
+    def _argument(self, arg: ast.arg, default: ast.expr | None, spaced: bool = False) -> str:
+        """`name=default` for main()'s inputs; `name = default` for a user function or an annotation."""
+        text = arg.arg
+        if arg.annotation is not None:
+            text += f': {self._expression(arg.annotation)}'
+            if default is not None:
+                text += f' = {self._expression(default)}'
+        elif default is not None:
+            text += f' = {self._expression(default)}' if spaced else f'={self._expression(default)}'
+        return text
+
+    def _arguments(self, args: ast.arguments, spaced: bool = False) -> list[str]:
+        positional = list(args.posonlyargs) + list(args.args)
+        defaults = [None] * (len(positional) - len(args.defaults)) + list(args.defaults)
+        parts = [self._argument(arg, default, spaced) for arg, default in zip(positional, defaults)]
+        if args.posonlyargs:
+            parts.insert(len(args.posonlyargs), '/')
+        if args.vararg is not None:
+            parts.append('*' + self._argument(args.vararg, None))
+        elif args.kwonlyargs:
+            parts.append('*')
+        parts.extend(self._argument(arg, default) for arg, default in zip(args.kwonlyargs, args.kw_defaults))
+        if args.kwarg is not None:
+            parts.append('**' + self._argument(args.kwarg, None))
+        return parts
+
+    def _function(self, node: ast.FunctionDef, depth: int, one_parameter_per_line: bool) -> list[str]:
+        indent = INDENT * depth
+        lines = [f'{indent}@{self._expression(decorator)}' for decorator in node.decorator_list]
+        keyword = 'async def' if isinstance(node, ast.AsyncFunctionDef) else 'def'
+        returns = f' -> {self._expression(node.returns)}' if node.returns is not None else ''
+        parameters = self._arguments(node.args, spaced=not one_parameter_per_line)
+        if one_parameter_per_line and parameters:
+            lines.append(f'{indent}{keyword} {node.name}(')
+            lines.append(',\n'.join(f'{indent}{INDENT}{parameter}' for parameter in parameters))
+            lines.append(f'{indent}){returns}:')
+        elif getattr(node, 'protocol_call', False) and parameters == ['self']:
+            lines.append(f'{indent}{keyword} {node.name}(self, ){returns}:')  # PyneComp's spelling
+        else:
+            lines.append(f'{indent}{keyword} {node.name}({", ".join(parameters)}){returns}:')
+        if _is_stub(node.body):
+            lines[-1] += ' ...'
+            return lines
+        return lines + self.block(node.body, depth + 1)
+
+    def _class(self, node: ast.ClassDef, depth: int) -> list[str]:
+        indent = INDENT * depth
+        lines = [f'{indent}@{self._expression(decorator)}' for decorator in node.decorator_list]
+        bases = [self._expression(base) for base in node.bases]
+        bases += [self._expression(keyword) for keyword in node.keywords]
+        head = f'{indent}class {node.name}({", ".join(bases)}):' if bases else f'{indent}class {node.name}:'
+        body: list[str] = []
+        for index, field in enumerate(node.body):
+            if index and self._blank_in_source(field):  # a Pine v6 UDT keeps the gaps between its fields
+                body.append('')
+            body += self.statement(field, depth + 1)
+        return lines + [head] + body
+
+    # -- statements --------------------------------------------------------
+
+    def _if(self, node: ast.If, depth: int, keyword: str = 'if') -> list[str]:
+        indent = INDENT * depth
+        lines = [f'{indent}{keyword} {self._expression(node.test)}:'] + self.block(node.body, depth + 1)
+        if node.orelse:
+            if len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If) and self._written_as_elif(node.orelse[0]):
+                lines += self._if(node.orelse[0], depth, 'elif')
+            else:
+                lines += [f'{indent}else:'] + self.block(node.orelse, depth + 1)
+        return lines
+
+    def _loop(self, node: ast.For | ast.While, depth: int) -> list[str]:
+        indent = INDENT * depth
+        if isinstance(node, ast.While):
+            head = f'{indent}while {self._expression(node.test)}:'
+        else:
+            target = node.target
+            spelled = (', '.join(self._expression(element) for element in target.elts)
+                       if isinstance(target, ast.Tuple) and target.elts else self._expression(target))
+            head = f'{indent}for {spelled} in {self._expression(node.iter)}:'
+        lines = [head] + self.block(node.body, depth + 1)
+        if node.orelse:
+            lines += [f'{indent}else:'] + self.block(node.orelse, depth + 1)
+        return lines
+
+    def statement(self, node: ast.stmt, depth: int, one_parameter_per_line: bool = False) -> list[str]:
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            return self._function(node, depth, one_parameter_per_line)
+        if isinstance(node, ast.ClassDef):
+            return self._class(node, depth)
+        if isinstance(node, ast.If):
+            return self._if(node, depth)
+        if isinstance(node, (ast.For, ast.While)):
+            return self._loop(node, depth)
+        indent = INDENT * depth
+        if isinstance(node, (ast.Try, ast.With, ast.Match, ast.AsyncFor, ast.AsyncWith)):
+            return [indent + line for line in ast.unparse(node).split('\n')]  # PyneComp does not emit these
+        return [indent + self._expression(node)]
+
+    def block(self, body: list[ast.stmt], depth: int) -> list[str]:
+        lines: list[str] = []
+        previous: ast.stmt | None = None
+        for node in body:
+            if previous is not None:
+                lines += [''] * self._blank_lines.blank_lines_before(previous, node)
+            lines += self.statement(node, depth)
+            previous = node
+        return lines
+
+
+def _is_stub(body: list[ast.stmt]) -> bool:
+    return (len(body) == 1 and isinstance(body[0], ast.Expr)
+            and isinstance(body[0].value, ast.Constant) and body[0].value.value is ...)
+
+
+def _is_docstring(node: ast.stmt) -> bool:
+    return isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant) and isinstance(node.value.value, str)
+
+
+def _is_main_guard(node: ast.stmt) -> bool:
+    return (isinstance(node, ast.If) and isinstance(node.test, ast.Compare)
+            and isinstance(node.test.left, ast.Name) and node.test.left.id == '__name__')
+
+
+def _binds_shadowed_namespace(node: ast.stmt) -> bool:
+    return (isinstance(node, ast.Assign) and isinstance(node.value, ast.Call)
+            and isinstance(node.value.func, ast.Name) and node.value.func.id == 'shadowed_namespace')
+
+
+def _declares_an_export(node: ast.stmt | None) -> bool:
+    """`name: _ProtocolName = Exported()`."""
+    value = getattr(node, 'value', None)
+    return isinstance(value, ast.Call) and isinstance(value.func, ast.Name) and value.func.id == 'Exported'
+
+
+def _assigned_names(node: ast.stmt) -> set[str]:
+    targets = node.targets if isinstance(node, ast.Assign) else [getattr(node, 'target', None)]
+    return {target.id for target in targets if isinstance(target, ast.Name)}
+
+
+def _moved_lines(source: LayoutSource, tree: ast.Module | None) -> list[int]:
+    """The Pine lines of the statements moved out of main(): the module-level assignments and the
+    enums and UDTs (found by their declaration in the Pine source)."""
+    if tree is None:
+        return []
+    lines = [source.origins.get(node.lineno, 0) for node in tree.body
+             if isinstance(node, (ast.Assign, ast.AnnAssign)) and source.origins.get(node.lineno, 0)]
+    for node in tree.body:
+        if isinstance(node, ast.ClassDef):
+            pine_name = re.sub(r'__ren__$', '', node.name)
+            match = re.search(rf'^(?:export\s+)?(?:type|enum)\s+{re.escape(pine_name)}\b', source.pine_source, re.M)
+            if match:
+                lines.append(source.pine_source.count('\n', 0, match.start()) + 1)
+    return lines
+
+
+class PyneCompModulePrinter:
+    """Prints a generated module in PyneComp's layout."""
+
+    def __init__(self, banner: PyneCompBanner | None = None, imports: ImportBlockFormatter | None = None) -> None:
+        self._banner = banner or PyneCompBanner()
+        self._imports = imports or ImportBlockFormatter()
+
+    def blank_line_policy(self, source: LayoutSource, tree: ast.Module | None = None) -> BlankLinePolicy:
+        if source.keeps_source_layout:
+            return SourceLayoutPolicy(source.origins, source.pine_source.split('\n'), _moved_lines(source, tree))
+        return AfterDefinitionPolicy()
+
+    _declares_types = False  # set per render: the module has a UDT or an enum
+
+    def _module_gap(self, previous: ast.stmt | None, current: ast.stmt) -> int:
+        """Blank lines between two module-level statements."""
+        if _binds_shadowed_namespace(current):
+            return 2 if previous is None or not _binds_shadowed_namespace(previous) else 0
+        if previous is not None and _binds_shadowed_namespace(previous) and not isinstance(current, _DEFINITIONS):
+            return 1
+        if previous is None:  # right after the import block
+            if isinstance(current, _DEFINITIONS) or '__all__' in _assigned_names(current):
+                return 2
+            return 1
+        if isinstance(previous, _DEFINITIONS) or isinstance(current, _DEFINITIONS):
+            return 2
+        if '__all__' in _assigned_names(previous) or '__all__' in _assigned_names(current):
+            return 2
+        if _declares_an_export(previous) != _declares_an_export(current):
+            # the Exported() declarations are one group; a constant after them is one line away, two in a
+            # library that declares a UDT or an enum (probes const2, ex1..ex8; 16 of 16 libraries)
+            if _declares_an_export(previous):
+                return 2 if self._declares_types else 1
+            return 2
+        return 0
+
+    def render(self, source: LayoutSource, tree: ast.Module | None = None,
+               spellings: Mapping[int, str] | None = None) -> str:
+        """Print ``tree`` (default: ``source.python_source`` parsed) in PyneComp's layout.
+
+        ``tree`` may be a rewritten parse of ``source.python_source``: nodes that
+        kept their positions still find their source spelling and Pine line.
+        """
+        if tree is None:
+            tree = ast.parse(source.python_source)
+        source_quoted = source_quoted_strings(tree) if source.keeps_source_layout else frozenset()
+        single_quoted = frozenset() if source.keeps_source_layout else header_strings(tree)
+        # the strings PyneComp writes as the source has them keep their non-printable characters raw
+        raw_strings = source_quoted if source.keeps_source_layout else source_quoted_strings(tree)
+        expression = ExpressionPrinter(source.python_source, source_quoted, spellings, single_quoted, raw_strings)
+        source_lines = source.python_source.split('\n')
+
+        def written_as_elif(node: ast.If) -> bool:
+            """Pine `else if` is `elif`; an `if` nested in an `else` block stays nested."""
+            lineno = getattr(node, 'lineno', 0)
+            return not lineno or source_lines[lineno - 1].lstrip().startswith('elif ')
+
+        def blank_in_source(node: ast.stmt) -> bool:
+            """The generated source has a blank line just above `node` (a gap between two UDT fields)."""
+            lineno = getattr(node, 'lineno', 0)
+            return lineno > 1 and not source_lines[lineno - 2].strip()
+
+        writer = _ModuleWriter(expression, self.blank_line_policy(source, tree), written_as_elif, blank_in_source)
+
+        self._declares_types = any(
+            isinstance(node, ast.ClassDef) and not any(isinstance(base, ast.Name) and base.id == 'Protocol'
+                                                       for base in node.bases) for node in tree.body)
+        body = [node for node in tree.body if not _is_docstring(node) and not _is_main_guard(node)]
+        import_nodes = [node for node in body if isinstance(node, (ast.Import, ast.ImportFrom))]
+        others = [node for node in body if not isinstance(node, (ast.Import, ast.ImportFrom))]
+
+        lines = self._banner.header() + self._imports.format(import_nodes)
+        previous: ast.stmt | None = None
+        for node in others:
+            lines += [''] * self._module_gap(previous, node)
+            is_main = isinstance(node, ast.FunctionDef) and node.name == 'main'
+            lines += writer.statement(node, 0, one_parameter_per_line=is_main)
+            previous = node
+        lines += ['', ''] + self._banner.footer()
+        return '\n'.join(lines) + '\n'