pine2pyne 0.2.4__py3-none-any.whl

pine2pyne/emit/pynecomp/passes/annotations.py

@@ -0,0 +1,422 @@
+"""Type annotations the way PyneComp writes them.
+
+* A name is a series (``Series``, ``PersistentSeries``) only when its scope
+  reads it with history (``x[1]``); a ``var`` otherwise is ``Persistent``.
+* A declaration (the first assignment of a name in a scope; later ones are
+  Pine ``:=``) is annotated when its type is known without inference: the Pine
+  source types it (``float x = ...``) or its value says it — a literal, a hex
+  colour, ``array.new_float()``, ``array.from_items('a', 'b')``, ``line.all``.
+* A series or persistent annotation keeps its element type (``Series[float]``)
+  only then; otherwise it is bare (``Series``).
+
+Measured on ``assets/pyne``: 2626 ``str``, 1254 ``int``, 1048 ``float``, 1146
+``bool``, 941 ``Color`` declarations; ``Series`` 2108 bare against 859 typed.
+"""
+from __future__ import annotations
+
+import ast
+import re
+
+from .base import PassContext, function_scopes, own_nodes
+
+_LITERAL_TYPES = {bool: 'bool', int: 'int', float: 'float', str: 'str'}
+_ARRAY_ELEMENTS = {
+    'new_float': 'float', 'new_int': 'int', 'new_bool': 'bool', 'new_string': 'str', 'new_color': 'Color',
+    'new_line': 'Line', 'new_label': 'Label', 'new_box': 'Box', 'new_table': 'Table',
+    'new_linefill': 'LineFill', 'new_polyline': 'Polyline',
+}
+_ALL_COLLECTIONS = {'line': 'Line', 'label': 'Label', 'box': 'Box', 'table': 'Table',
+                    'linefill': 'LineFill', 'polyline': 'Polyline'}
+_SERIES_KINDS = {'Series': None, 'PersistentSeries': 'Persistent', 'IBPersistentSeries': 'IBPersistent'}
+_PERSISTENT_KINDS = ('Series', 'PersistentSeries', 'IBPersistentSeries', 'Persistent', 'IBPersistent')
+
+
+def _dotted(node: ast.AST) -> str:
+    if isinstance(node, ast.Attribute):
+        return f'{_dotted(node.value)}.{node.attr}'
+    return node.id if isinstance(node, ast.Name) else ''
+
+
+def literal_type(value: ast.expr | None) -> str | None:
+    if isinstance(value, ast.Constant):
+        return _LITERAL_TYPES.get(type(value.value))
+    if (isinstance(value, ast.UnaryOp) and isinstance(value.op, ast.USub)
+            and isinstance(value.operand, ast.Constant) and type(value.operand.value) is int):
+        return 'int'
+    return None
+
+
+def _typed_na_argument(call: ast.Call) -> str | None:
+    """T of the `NA(T)` value of `array.new(n, NA(T))` / `matrix.new(..., initial_value=NA(T))`."""
+    position = 2 if _dotted(call.func) == 'matrix.new' else 1  # `matrix.new(r, c, NA(T))`
+    values = list(call.args[position:position + 1]) + [keyword.value for keyword in call.keywords
+                                                        if keyword.arg == 'initial_value']
+    for value in values:
+        if (isinstance(value, ast.Call) and isinstance(value.func, ast.Name) and value.func.id == 'NA'
+                and len(value.args) == 1 and isinstance(value.args[0], ast.Name)):
+            return value.args[0].id
+    return None
+
+
+def inferable_type(value: ast.expr | None, classes: frozenset = frozenset(),
+                   methods: frozenset = frozenset()) -> str | None:
+    """The type PyneComp reads off a value without inference, or None. ``classes``:
+    the module's UDT classes, whose constructor call names its type; ``methods``: the user
+    methods, whose call on a constructor (`T.new().m()`) PyneComp types as `T` (probe mc5)."""
+    found = literal_type(value)
+    if found:
+        return found
+    if isinstance(value, ast.Attribute) and value.attr == 'all' and isinstance(value.value, ast.Name):
+        element = _ALL_COLLECTIONS.get(value.value.id)
+        return f'list[{element}]' if element else None
+    if not isinstance(value, ast.Call):
+        return None
+    function = _dotted(value.func)
+    if function in classes:
+        return function
+    if function.endswith('.new') and function.count('.') == 2 and function.split('.')[1][:1].isupper():
+        return function[:-len('.new')]  # a library's type, `dr.AuxData.new(...)` (biased_price_target)
+    receiver = None
+    if function == 'method_call' and len(value.args) >= 2:
+        receiver = value.args[1]
+    elif function in methods and value.args:
+        receiver = value.args[0]
+    if isinstance(receiver, ast.Call):
+        return inferable_type(receiver, classes, methods) if _dotted(receiver.func) in classes | methods \
+            or _dotted(receiver.func) == 'method_call' else None
+    if function == 'color.new' and len(value.args) == 1 and not value.keywords \
+            and isinstance(value.args[0], ast.Constant) and isinstance(value.args[0].value, str):
+        return 'Color'  # a hex colour literal
+    if function == 'matrix.new' and len(value.args) >= 3 and literal_type(value.args[2]):
+        return f'Matrix[{literal_type(value.args[2])}]'  # the element type of the initial value
+    typed_na = _typed_na_argument(value)
+    if function in ('matrix.new', 'array.new') and typed_na:
+        return f"{'Matrix' if function == 'matrix.new' else 'list'}[{typed_na}]"  # `NA(T)`: of T
+    if function.startswith('array.'):
+        element = _ARRAY_ELEMENTS.get(function[len('array.'):])
+        if element:
+            return f'list[{element}]'
+        if function == 'array.from_items' and value.args:
+            types = {literal_type(argument) for argument in value.args}
+            if len(types) == 1 and None not in types:
+                return f'list[{types.pop()}]'
+    return None
+
+
+class DeclaredTypes:
+    """The Python types of the Pine source's typed declarations (``float x = ...``),
+    looked up by name and, when the statement's Pine line is known, by that line."""
+
+    def __init__(self, script) -> None:
+        from ....ast_nodes import Assignment, VarDecl, VaripDecl
+        from ....ast_walk import walk
+        from ....transformer import convert_pine_generic_to_python
+        self.by_line: dict = {}
+        self.by_name: dict = {}
+        self.declared_lines: dict = {}
+        for node in walk([getattr(script, 'declarations', []), getattr(script, 'body', [])]):
+            if isinstance(node, (Assignment, VarDecl, VaripDecl)):
+                name = node.target if isinstance(node, Assignment) else node.name
+                if not isinstance(name, str):
+                    continue
+                line = getattr(node, 'line', 0)
+                self.declared_lines.setdefault(name, set()).add(line)
+                type_hint = getattr(node, 'type_hint', None) or self._generic_constructor_type(node.value)
+                if type_hint:
+                    python_type = convert_pine_generic_to_python(type_hint)
+                    self.by_line[(name, line)] = python_type
+                    self.by_name.setdefault(name, python_type)
+
+    @staticmethod
+    def _generic_constructor_type(value) -> str | None:
+        """`map.new<string, bool>()` names its type, `map<string, bool>` (probe nb4_02); `T.copy(x)` names `T`
+        (Trendoscope/Zigzag `lPivot: Pivot = udt_copy(...)`)."""
+        from ....ast_nodes import FunctionCall
+        if isinstance(value, FunctionCall) and isinstance(value.func, str) and value.func.startswith('map.new<'):
+            return 'map' + value.func[len('map.new'):]
+        if isinstance(value, FunctionCall) and isinstance(value.func, str) and value.func.endswith('.copy') \
+                and value.func.count('.') == 1 and value.func[:1].isupper() and len(value.args) == 1:
+            return value.func[:-len('.copy')]
+        return None
+
+    #: the suffixes renaming gives a Pine name: `x__000000c0__`, `x__ren__`, `x__global__`
+    _RENAMED = re.compile(r'__[0-9a-f]{8}__$|__ren__$|__global__$')
+
+    def get(self, name: str, line: int = 0):
+        name = self._RENAMED.sub('', name)
+        if line and line in self.declared_lines.get(name, ()):
+            return self.by_line.get((name, line))
+        return self.by_name.get(name)
+
+    def __contains__(self, name: str) -> bool:
+        return self._RENAMED.sub('', name) in self.by_name
+
+
+def pine_declared_types(script) -> DeclaredTypes:
+    return DeclaredTypes(script)
+
+
+class SeriesAnnotationPass:
+    """A name is a series only where its scope reads it with history: otherwise
+    ``PersistentSeries[T]`` is ``Persistent[T]`` and a ``Series`` annotation goes."""
+
+    @staticmethod
+    def _history_reads(scope: ast.AST) -> set[str]:
+        """Names `scope` reads with history, its nested functions included where they do
+        not bind the name themselves (a parameter `min` is not main's `min`)."""
+        reads = {node.value.id for node in own_nodes(scope)
+                 if isinstance(node, ast.Subscript) and isinstance(node.value, ast.Name)}
+        for node in own_nodes(scope):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                arguments = node.args
+                bound = {arg.arg for arg in arguments.posonlyargs + arguments.args + arguments.kwonlyargs}
+                bound |= {child.id for child in own_nodes(node)
+                          if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Store)}
+                reads |= SeriesAnnotationPass._history_reads(node) - bound
+        return reads
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        for scope in function_scopes(tree):
+            read_with_history = self._history_reads(scope)
+            blocks = [node for node in own_nodes(scope) if hasattr(node, 'body')
+                      and not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))]
+            for owner in [scope] + blocks:
+                for field in ('body', 'orelse', 'finalbody'):
+                    block = getattr(owner, field, None)
+                    if isinstance(block, list):
+                        block[:] = [self._demoted(statement, read_with_history) for statement in block]
+
+    @staticmethod
+    def _demoted(statement: ast.stmt, read_with_history: set) -> ast.stmt:
+        if not (isinstance(statement, ast.AnnAssign) and isinstance(statement.target, ast.Name)
+                and statement.target.id not in read_with_history):
+            return statement
+        annotation = statement.annotation
+        head = annotation.value if isinstance(annotation, ast.Subscript) else annotation
+        if not isinstance(head, ast.Name) or head.id not in _SERIES_KINDS:
+            return statement
+        replacement = _SERIES_KINDS[head.id]
+        if replacement is None:  # a plain series: no annotation
+            if statement.value is None:
+                return statement
+            return ast.copy_location(ast.Assign(targets=[statement.target], value=statement.value), statement)
+        if isinstance(annotation, ast.Subscript):
+            annotation.value = ast.Name(id=replacement, ctx=ast.Load())
+        else:
+            statement.annotation = ast.Name(id=replacement, ctx=ast.Load())
+        return statement
+
+
+def _unquoted(annotation) -> str:
+    """An annotation's text, a string annotation (`'list[Pivot]'`) unquoted."""
+    if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str):
+        return annotation.value
+    return ast.unparse(annotation)
+
+
+def _element_of_kind(annotation: str | None) -> str | None:
+    """`Persistent[Pivot]`, `Series[Pivot]` -> `Pivot`."""
+    if annotation is None:
+        return None
+    match = re.fullmatch(r'(?:Persistent|PersistentSeries|Series|IBPersistent|IBPersistentSeries)\[(.*)\]', annotation)
+    return match.group(1) if match else annotation
+
+
+class DeclarationAnnotationPass:
+    """Declarations carry a type PyneComp knows; series and persistent annotations
+    keep their element type only then (see the module docstring)."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        typed = pine_declared_types(context.unit.pine_script)
+        self._origins = context.origins
+        self._classes = frozenset(node.name for node in tree.body if isinstance(node, ast.ClassDef))
+        self._methods = frozenset(node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)
+                                  and any(isinstance(d, ast.Name) and d.id == 'method' for d in node.decorator_list))
+        # each UDT's fields and their annotations: a field read is typed by its field (probe nodes1)
+        self._fields = {node.name: {field.target.id: _unquoted(field.annotation) for field in node.body
+                                    if isinstance(field, ast.AnnAssign) and isinstance(field.target, ast.Name)}
+                        for node in tree.body if isinstance(node, ast.ClassDef)}
+        for scope in function_scopes(tree):
+            declared: set[str] = set()
+            types: dict[str, str] = {}
+            if isinstance(scope, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                arguments = scope.args
+                for arg in arguments.posonlyargs + arguments.args + arguments.kwonlyargs:
+                    declared.add(arg.arg)
+                    if arg.annotation is not None:
+                        types[arg.arg] = _unquoted(arg.annotation)
+            self._annotate(scope.body, declared, typed, types)
+
+    _classes: frozenset = frozenset()
+    _methods: frozenset = frozenset()
+    _fields: dict = {}
+    _origins: dict = {}
+
+    def _known(self, name: str, value, typed, statement=None) -> str | None:
+        line = self._origins.get(getattr(statement, 'lineno', 0), 0) if statement is not None else 0
+        return typed.get(name, line) or inferable_type(value, self._classes, self._methods)
+
+    def _bare_unless_known(self, statement: ast.AnnAssign, typed: dict[str, str],
+                           types: dict[str, str] | None = None) -> None:
+        annotation = statement.annotation
+        if isinstance(annotation, ast.Name) and annotation.id in _PERSISTENT_KINDS:
+            # a bare kind takes the type of a constructor's method chain (`var w = T.new().init()`), of a
+            # library type's constructor (`var s = lib.T.new()`) (chrono_utils, biased_price_target) or of the
+            # field it reads (`cl: Series[float] = c.c`, Trendoscope/oota)
+            element = self._constructed_type(statement.value) or self._field_type(statement.value, typed, types or {})
+            if element:
+                statement.annotation = ast.Subscript(value=annotation, slice=ast.parse(element, mode='eval').body,
+                                                     ctx=ast.Load())
+            return
+        if not (isinstance(annotation, ast.Subscript) and isinstance(annotation.value, ast.Name)
+                and annotation.value.id in _PERSISTENT_KINDS):
+            return
+        if not (self._known(statement.target.id, statement.value, typed, statement)
+                or self._field_type(statement.value, typed, types or {})):
+            statement.annotation = ast.Name(id=annotation.value.id, ctx=ast.Load())
+        elif literal_type(statement.value) and statement.target.id not in typed:
+            # a literal names the element type itself: `x: Series[float] = 0.0`
+            annotation.slice = ast.Name(id=literal_type(statement.value), ctx=ast.Load())
+
+    def _constructed_type(self, value) -> str | None:
+        if isinstance(value, ast.Call):
+            function = _dotted(value.func)
+            if function.endswith('.new') and function.count('.') == 2 and function.split('.')[1][:1].isupper():
+                return function[:-len('.new')]  # `dr.DistSettings.new(...)`
+            if function in self._methods or function == 'method_call':
+                return inferable_type(value, self._classes, self._methods)
+        return None
+
+    def _field_type(self, value, typed, types: dict[str, str]) -> str | None:
+        """The declared type of the UDT field `value` reads (`p.x`, `a.b.c`), when its receiver's type is known."""
+        if not isinstance(value, ast.Attribute):
+            return None
+        receiver = value.value
+        if isinstance(receiver, ast.Name):
+            owner = types.get(receiver.id)  # only a type this scope declares (Zigzag `llastPivot` is untyped)
+        else:
+            owner = self._field_type(receiver, typed, types)
+        owner = _element_of_kind(owner)
+        return self._fields.get(owner, {}).get(value.attr) if owner else None
+
+    def _annotate(self, body: list[ast.stmt], declared: set[str], typed: dict[str, str],
+                  types: dict[str, str] | None = None) -> None:
+        types = {} if types is None else types
+        for index, statement in enumerate(body):
+            if isinstance(statement, ast.Assign) and len(statement.targets) == 1 \
+                    and isinstance(statement.targets[0], ast.Name):
+                name = statement.targets[0].id
+                value = statement.value
+                if isinstance(value, ast.Attribute) and isinstance(value.value, ast.Name) \
+                        and re.fullmatch(r'__loop_\d+__', value.value.id):
+                    declared.add(name)  # a `pine_loop` variable's value is never annotated (analytics_tables)
+                    continue
+                annotation = self._known(name, statement.value, typed, statement)
+                if annotation is None and name not in declared:
+                    annotation = self._field_type(statement.value, typed, types)
+                if name not in declared and annotation and not name.startswith('__'):
+                    body[index] = ast.copy_location(ast.AnnAssign(
+                        target=statement.targets[0], annotation=ast.parse(annotation, mode='eval').body,
+                        value=statement.value, simple=1), statement)
+                    types[name] = annotation
+                declared.add(name)
+            elif isinstance(statement, ast.Assign):  # `a, b = f()` declares its names
+                declared.update(node.id for target in statement.targets for node in ast.walk(target)
+                                if isinstance(node, ast.Name))
+            elif isinstance(statement, ast.AnnAssign) and isinstance(statement.target, ast.Name):
+                self._bare_unless_known(statement, typed, types)
+                declared.add(statement.target.id)
+                types[statement.target.id] = _unquoted(statement.annotation)
+            elif isinstance(statement, (ast.For, ast.AsyncFor)):
+                declared.update(node.id for node in ast.walk(statement.target) if isinstance(node, ast.Name))
+            elif isinstance(statement, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                declared.add(statement.name)
+                continue
+            for name in ('body', 'orelse', 'finalbody'):
+                nested = getattr(statement, name, None)
+                if isinstance(nested, list):
+                    # a Pine block is a scope: what it declares is unknown to its siblings
+                    self._annotate(nested, set(declared), typed, dict(types))
+
+
+class ColorCastPass:
+    """A colour value is not wrapped in ``Color(...)``: ``GREEN: Color = color.new('#00FF00')``."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        for node in ast.walk(tree):
+            for field, value in ast.iter_fields(node):
+                if (isinstance(value, ast.Call) and isinstance(value.func, ast.Name) and value.func.id == 'Color'
+                        and len(value.args) == 1 and not value.keywords and isinstance(value.args[0], ast.Call)):
+                    setattr(node, field, value.args[0])
+
+
+class ClosureHistoryPass:
+    """A function that reads the history of a variable of its enclosing function
+    (``avg[1]`` with ``avg`` a main() variable) reads it from its own series,
+    ``__hist_avg__: Series = avg[0]``, bound first in its body (33 modules of the reference)."""
+
+    PREFIX, SUFFIX = '__hist_', '__'
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        main = next((node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == 'main'), None)
+        if main is not None:
+            # the snapshot keeps the element type the variable is declared with in main():
+            # `x: PersistentSeries[float]` or `x: float` gives `__hist_x__: Series[float]`
+            self._series_types = {}
+            declared = [(arg.arg, arg.annotation) for arg in main.args.args]
+            declared += [(node.target.id, node.annotation) for node in own_nodes(main)
+                         if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name)]
+            for name, annotation in declared:
+                element = self._element_type(annotation)
+                if element is not None and name not in self._series_types:
+                    self._series_types[name] = ast.Subscript(value=ast.Name(id='Series', ctx=ast.Load()),
+                                                             slice=element, ctx=ast.Load())
+            self._visit(main, set())
+
+    _series_types: dict = {}
+
+    @staticmethod
+    def _element_type(annotation):
+        if annotation is None:
+            return None
+        if isinstance(annotation, ast.Subscript) and isinstance(annotation.value, ast.Name) \
+                and annotation.value.id in _PERSISTENT_KINDS:
+            return annotation.slice
+        if isinstance(annotation, ast.Name) and annotation.id not in _PERSISTENT_KINDS:
+            return annotation
+        return None
+
+    @staticmethod
+    def _own_bindings(function: ast.FunctionDef) -> set[str]:
+        arguments = function.args
+        bound = {arg.arg for arg in arguments.posonlyargs + arguments.args + arguments.kwonlyargs}
+        for node in own_nodes(function):
+            if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store):
+                bound.add(node.id)
+            elif isinstance(node, (ast.FunctionDef, ast.ClassDef)):
+                bound.add(node.name)
+        return bound
+
+    def _visit(self, function: ast.FunctionDef, enclosing: set[str]) -> None:
+        own = self._own_bindings(function)
+        if enclosing:
+            snapshots: list[str] = []
+            for node in ast.walk(function):
+                if (isinstance(node, ast.Subscript) and isinstance(node.value, ast.Name)
+                        and node.value.id in enclosing and node.value.id not in own
+                        and node.value.id not in snapshots):
+                    snapshots.append(node.value.id)
+            if snapshots:
+                for node in ast.walk(function):
+                    if (isinstance(node, ast.Subscript) and isinstance(node.value, ast.Name)
+                            and node.value.id in snapshots):
+                        node.value = ast.Name(id=f'{self.PREFIX}{node.value.id}{self.SUFFIX}', ctx=ast.Load())
+                first = function.body[0] if function.body else function
+                function.body[:0] = [ast.copy_location(ast.AnnAssign(
+                    target=ast.Name(id=f'{self.PREFIX}{name}{self.SUFFIX}', ctx=ast.Store()),
+                    annotation=self._series_types.get(name) or ast.Name(id='Series', ctx=ast.Load()),
+                    value=ast.Subscript(value=ast.Name(id=name, ctx=ast.Load()), slice=ast.Constant(0),
+                                        ctx=ast.Load()), simple=1), first) for name in snapshots]
+        for node in own_nodes(function):
+            if isinstance(node, ast.FunctionDef):
+                self._visit(node, (enclosing | own) - self._own_bindings(node))

pine2pyne/emit/pynecomp/passes/base.py

@@ -0,0 +1,69 @@
+"""Rewrites of the generated Python module towards PyneComp's output.
+
+pine2pyne and PyneComp agree on what a script computes but often spell it
+differently (``PersistentSeries`` vs ``Persistent``, ``'#00ced1'`` vs
+``color.new('#00CED1')``). Each difference is one :class:`ModulePass`: a
+small rewrite of the Python AST, applied in order by the PyneComp emitter.
+A new difference is a new pass; existing passes stay untouched.
+"""
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass, field
+from typing import Iterator, Protocol
+
+from ....core.contracts import CompilationUnit
+
+
+@dataclass
+class PassContext:
+    unit: CompilationUnit
+    #: exact text to print for a node (``id(node)``), e.g. a title's source quotes
+    spellings: dict[int, str] = field(default_factory=dict)
+    #: Python line -> the Pine line of the statement generated there
+    origins: dict[int, int] = field(default_factory=dict)
+
+    @property
+    def pine_version(self) -> int:
+        return self.unit.pine_version or 6
+
+
+class ModulePass(Protocol):
+    def apply(self, tree: ast.Module, context: PassContext) -> None: ...
+
+
+def main_function(tree: ast.Module) -> ast.FunctionDef | None:
+    for node in tree.body:
+        if isinstance(node, ast.FunctionDef) and node.name == 'main':
+            return node
+    return None
+
+
+def function_scopes(tree: ast.Module) -> Iterator[ast.AST]:
+    """The module and every function in it."""
+    yield tree
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            yield node
+
+
+def own_nodes(scope: ast.AST) -> Iterator[ast.AST]:
+    """Nodes of ``scope`` that are not inside a nested function or class."""
+    stack = list(ast.iter_child_nodes(scope))
+    while stack:
+        node = stack.pop()
+        yield node
+        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+            stack.extend(ast.iter_child_nodes(node))
+
+
+def statements_in_order(body: list[ast.stmt]) -> Iterator[ast.stmt]:
+    """Statements of a block and its nested blocks (not nested functions), in source order."""
+    for statement in body:
+        yield statement
+        if isinstance(statement, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+            continue
+        for name in ('body', 'orelse', 'finalbody'):
+            nested = getattr(statement, name, None)
+            if isinstance(nested, list):
+                yield from statements_in_order(nested)

pine2pyne/emit/pynecomp/passes/builtins.py

@@ -0,0 +1,169 @@
+"""Built-in calls spelled the way PyneComp spells them."""
+from __future__ import annotations
+
+import ast
+
+from .base import PassContext
+
+#: ``ta`` members that are series values, not calls (``ta.tr(true)`` is the call form)
+TA_SERIES = frozenset({'tr', 'vwap', 'obv', 'accdist', 'iii', 'nvi', 'pvi', 'pvt', 'wad', 'wvad'})
+#: Pine v4 hline and plot styles were integers
+_HLINE_STYLES = {0: 'style_solid', 1: 'style_dotted', 2: 'style_dashed'}
+_PLOT_STYLES = {0: 'style_line', 1: 'style_stepline', 2: 'style_histogram', 3: 'style_cross', 4: 'style_area',
+                5: 'style_columns', 6: 'style_circles'}
+
+
+def _dotted(node: ast.AST) -> str:
+    if isinstance(node, ast.Attribute):
+        return f'{_dotted(node.value)}.{node.attr}'
+    return node.id if isinstance(node, ast.Name) else ''
+
+
+def _attribute(base: str, member: str) -> ast.Attribute:
+    return ast.Attribute(value=ast.Name(id=base, ctx=ast.Load()), attr=member, ctx=ast.Load())
+
+
+def _bound_names(tree: ast.Module) -> set[str]:
+    bound = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store)}
+    bound |= {node.arg for node in ast.walk(tree) if isinstance(node, ast.arg)}
+    bound |= {node.name for node in ast.walk(tree) if isinstance(node, (ast.FunctionDef, ast.ClassDef))}
+    return bound
+
+
+class _Replacer(ast.NodeTransformer):
+    def __init__(self, replace):
+        self._replace = replace
+
+    def generic_visit(self, node):
+        node = super().generic_visit(node)
+        replaced = self._replace(node)
+        return node if replaced is None else ast.copy_location(replaced, node)
+
+
+class TaSeriesPass:
+    """``ta.tr``, ``ta.vwap``, ``ta.obv``, ... are values: no call without arguments;
+    a Pine v4 bare ``tr``/``vwap``/... is the ``ta.`` value."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        bound = _bound_names(tree)
+        legacy = context.pine_version <= 4
+
+        def replace(node):
+            if (isinstance(node, ast.Call) and not node.args and not node.keywords
+                    and isinstance(node.func, ast.Attribute) and node.func.attr in TA_SERIES
+                    and _dotted(node.func.value) == 'ta'):
+                return node.func
+            if (legacy and isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load)
+                    and node.id in TA_SERIES and node.id not in bound):
+                return _attribute('ta', node.id)
+            return None
+        _Replacer(replace).visit(tree)
+
+
+class LegacyStrategyDirectionPass:
+    """Pine v4 ``strategy.entry(id, true)`` / ``long=true``: the direction is
+    ``strategy.long`` / ``strategy.short`` (``direction=`` for the keyword)."""
+
+    FUNCTIONS = frozenset({'strategy.entry', 'strategy.order'})
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        if context.pine_version > 4:
+            return
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Call) and _dotted(node.func) in self.FUNCTIONS:
+                if len(node.args) >= 2:
+                    node.args[1] = self._direction(node.args[1])
+                for keyword in node.keywords:
+                    if keyword.arg == 'long':
+                        keyword.arg, keyword.value = 'direction', self._direction(keyword.value)
+
+    @staticmethod
+    def _direction(value: ast.expr) -> ast.expr:
+        if isinstance(value, ast.Constant) and value.value in (True, False, 1, 0):
+            return _attribute('strategy', 'long' if value.value else 'short')
+        return value
+
+
+class LegacyLineWidthPass:
+    """Pine v4/v5: a ``linewidth`` that is not a constant of at least 1 is ``math.max(1, width)``."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        if context.pine_version >= 6:
+            return
+        for node in ast.walk(tree):
+            at_least_one = (isinstance(getattr(node, 'value', None), ast.Constant)
+                            and isinstance(node.value.value, (int, float)) and node.value.value >= 1)
+            if isinstance(node, ast.keyword) and node.arg == 'linewidth' \
+                    and not at_least_one and not self._is_max(node.value):
+                node.value = ast.Call(func=_attribute('math', 'max'), args=[ast.Constant(1), node.value], keywords=[])
+
+    @staticmethod
+    def _is_max(value: ast.expr) -> bool:
+        return isinstance(value, ast.Call) and _dotted(value.func) == 'math.max'
+
+
+class CurrentValuePass:
+    """``x[0]`` is ``x``."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        def replace(node):
+            if (isinstance(node, ast.Subscript) and isinstance(node.ctx, ast.Load)
+                    and isinstance(node.slice, ast.Constant) and node.slice.value == 0
+                    and isinstance(node.value, (ast.Name, ast.Attribute))
+                    and not (isinstance(node.value, ast.Name) and node.value.id.startswith('__'))):
+                return node.value
+            return None
+        _Replacer(replace).visit(tree)
+
+
+class LegacyHlineStylePass:
+    """Pine v4 ``hline(..., linestyle=2)`` / ``plot(..., style=5)``: the style constant,
+    ``hline.style_dashed`` / ``plot.style_columns``."""
+
+    _STYLES = {('hline', 'linestyle'): ('hline', _HLINE_STYLES), ('plot', 'style'): ('plot', _PLOT_STYLES)}
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        if context.pine_version > 4:
+            return
+        for node in ast.walk(tree):
+            if not isinstance(node, ast.Call):
+                continue
+            for keyword in node.keywords:
+                namespace, styles = self._STYLES.get((_dotted(node.func), keyword.arg), (None, {}))
+                if (namespace and isinstance(keyword.value, ast.Constant)
+                        and type(keyword.value.value) is int and keyword.value.value in styles):
+                    keyword.value = _attribute(namespace, styles[keyword.value.value])
+
+
+class InputOptionsPass:
+    """``options=`` is always a tuple: ``options=x`` (one name) is ``options=(x,)``."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Call) and _dotted(node.func).startswith('input'):
+                for keyword in node.keywords:
+                    if keyword.arg == 'options' and not isinstance(keyword.value, (ast.Tuple, ast.List)):
+                        keyword.value = ast.Tuple(elts=[keyword.value], ctx=ast.Load())
+
+
+class ShadowedNamespacePass:
+    """A library imported under a built-in namespace's name (``import TradingView/ta/12``)
+    reads through ``ta__ren__ = shadowed_namespace(ta__ren____lib, ta)``: the library's
+    members first, the built-in's after. Every ``ta.`` of the script reads ``ta__ren__``."""
+
+    def apply(self, tree: ast.Module, context: PassContext) -> None:
+        shadowed = getattr(context.unit.transformed, 'shadowed_namespaces', None) or {}
+        if not shadowed:
+            return
+        for node in ast.walk(tree):
+            if (isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name)
+                    and node.value.id in shadowed):
+                node.value = ast.Name(id=f'{node.value.id}__ren__', ctx=ast.Load())
+        position = max((index + 1 for index, node in enumerate(tree.body)
+                        if isinstance(node, (ast.Import, ast.ImportFrom))), default=0)
+        tree.body[position:position] = [
+            ast.Assign(targets=[ast.Name(id=f'{alias}__ren__', ctx=ast.Store())],
+                       value=ast.Call(func=ast.Name(id='shadowed_namespace', ctx=ast.Load()),
+                                      args=[ast.Name(id=module, ctx=ast.Load()), ast.Name(id=alias, ctx=ast.Load())],
+                                      keywords=[]), lineno=0, col_offset=0)
+            for alias, module in shadowed.items()]