@pineforge/codegen-pyodide 0.8.0 → 0.9.0

package/pineforge_codegen/analyzer/call_handlers.py

@@ -66,8 +66,8 @@ from __future__ import annotations
 from typing import Any
 
 from ..ast_nodes import (
-    ASTNode, BoolLiteral, FuncCall, Identifier, MemberAccess,
-    NumberLiteral, StringLiteral, TupleLiteral,
+    ASTNode, BinOp, BoolLiteral, ExprStmt, FuncCall, Identifier, MemberAccess,
+    NumberLiteral, StringLiteral, Ternary, TupleLiteral, UnaryOp, VarDecl,
 )
 from ..symbols import PineType
 from .. import signatures as sigs
@@ -75,7 +75,7 @@ from .. import tv_input_choices as tv_in
 from .contracts import FixnanCallSite, FuncInfo, SecurityCallInfo, TACallSite
 from .tables import (
     BAR_FIELDS, TA_CLASS_MAP, TA_MULTI_CTOR, TA_NO_CTOR, TA_PERIOD_ARG,
-    TA_TUPLE_RETURNS,
+    TA_TUPLE_RETURNS, TA_TUPLE_ELEMENT_COUNTS, TA_COMPUTE_ARGS,
 )
 
 
@@ -205,8 +205,10 @@ class CallHandlers:
                 returns_tuple=returns_tuple,
                 node=node,
                 is_static=is_static,
+                owner_func=(self._enclosing_func_names[-1] if self._enclosing_func_names else None),
             )
             self._ta_call_sites.append(site)
+            self._ta_member_names.add(site.member_name)
             return PineType.FLOAT
 
         # Determine constructor args
@@ -234,9 +236,14 @@ class CallHandlers:
         elif func_name in TA_PERIOD_ARG:
             ctor_indices = {TA_PERIOD_ARG[func_name]}
 
-        for i, arg in enumerate(all_args):
-            if i not in ctor_indices and arg is not None:
-                compute_args.append(arg)
+        if func_name in TA_COMPUTE_ARGS:
+            for i in TA_COMPUTE_ARGS[func_name]:
+                if i < len(all_args) and all_args[i] is not None:
+                    compute_args.append(all_args[i])
+        else:
+            for i, arg in enumerate(all_args):
+                if i not in ctor_indices and arg is not None:
+                    compute_args.append(arg)
 
         is_static = self._global_scope and all(self._is_static_expression(arg) for arg in compute_args)
         site = TACallSite(
@@ -247,11 +254,32 @@ class CallHandlers:
             returns_tuple=returns_tuple,
             node=node,
             is_static=is_static,
+            owner_func=(self._enclosing_func_names[-1] if self._enclosing_func_names else None),
         )
         self._ta_call_sites.append(site)
+        self._ta_member_names.add(site.member_name)
 
         return PineType.FLOAT
 
+    def _security_symbol_is_heikinashi(self, node, _seen=None) -> bool:
+        """True when a request.security symbol is ``ticker.heikinashi(<chart
+        symbol>)`` — directly, or via a global alias (``haTicker =
+        ticker.heikinashi(syminfo.tickerid)``). Name-cycle-guarded. The
+        support_checker has already rejected the cross-symbol HA case, so any HA
+        reaching here is the chart's own symbol."""
+        if _seen is None:
+            _seen = set()
+        if (isinstance(node, FuncCall) and isinstance(node.callee, MemberAccess)
+                and isinstance(node.callee.object, Identifier)
+                and node.callee.object.name == "ticker"
+                and node.callee.member == "heikinashi"):
+            return True
+        if (isinstance(node, Identifier) and node.name in self._global_expr_map
+                and node.name not in _seen):
+            _seen.add(node.name)
+            return self._security_symbol_is_heikinashi(self._global_expr_map[node.name], _seen)
+        return False
+
     def _handle_request_call(self, func_name: str, node: FuncCall) -> PineType:
         """Handle request.* function calls."""
         if func_name == "security":
@@ -288,11 +316,41 @@ class CallHandlers:
 
             returns_tuple = isinstance(expr_node, TupleLiteral)
             tuple_size = len(expr_node.elements) if returns_tuple else 0
+            if not returns_tuple and isinstance(expr_node, FuncCall):
+                expr_func = None
+                expr_ns = None
+                if (isinstance(expr_node.callee, MemberAccess)
+                        and isinstance(expr_node.callee.object, Identifier)):
+                    expr_ns = expr_node.callee.object.name
+                    expr_func = expr_node.callee.member
+                if expr_ns == "ta":
+                    if expr_func == "vwap":
+                        merged_v = list(expr_node.args)
+                        for i, pname in enumerate(["source", "anchor", "stdev_mult"]):
+                            if pname in expr_node.kwargs:
+                                while len(merged_v) <= i:
+                                    merged_v.append(None)
+                                if merged_v[i] is None:
+                                    merged_v[i] = expr_node.kwargs[pname]
+                        if len(merged_v) >= 3:
+                            expr_func = "vwap_bands"
+                    if expr_func in TA_TUPLE_RETURNS:
+                        returns_tuple = True
+                        tuple_size = TA_TUPLE_ELEMENT_COUNTS.get(expr_func, 0)
 
             gaps_node = all_args[3] if len(all_args) > 3 else None
             lookahead_node = all_args[4] if len(all_args) > 4 else None
 
             mutable_globals = tuple(sorted(self._collect_security_mutable_globals(expr_node)))
+            # Heikin-Ashi same-symbol read: request.security(ticker.heikinashi(
+            # syminfo.tickerid), ...) (directly or via a global alias). The engine
+            # applies the HA candle transform inside the security eval.
+            symbol_node = all_args[0] if all_args else None
+            heikinashi = self._security_symbol_is_heikinashi(symbol_node)
+            # Capture the user function (if any) whose body contains this call,
+            # so the codegen can resolve a parameter ``tf`` via the call sites.
+            scope_name = self._symbols.current_scope.name
+            containing_func = scope_name[5:] if scope_name.startswith("func_") else ""
             self._security_calls.append(SecurityCallInfo(
                 sec_id=sec_id,
                 timeframe=tf_node,
@@ -302,8 +360,10 @@ class CallHandlers:
                 gaps=gaps_node,
                 lookahead=lookahead_node,
                 ta_range=security_ta_range,
+                heikinashi=heikinashi,
                 depends_on_mutable_globals=bool(mutable_globals),
                 mutable_globals=mutable_globals,
+                containing_func=containing_func,
             ))
 
             return PineType.FLOAT
@@ -753,11 +813,18 @@ class CallHandlers:
             arg_type = self._visit(arg)
 
         self._fixnan_counter += 1
+        owner = self._enclosing_func_names[-1] if self._enclosing_func_names else None
         site = FixnanCallSite(
             member_name=f"_prev_fixnan_{self._fixnan_counter}",
             pine_type=arg_type,
+            node=node,
+            owner_func=owner,
         )
+        idx = len(self._fixnan_sites)
         self._fixnan_sites.append(site)
+        self._fixnan_member_names.add(site.member_name)
+        if owner is not None:
+            self._func_fixnan_indices.setdefault(owner, []).append(idx)
 
         return arg_type
 
@@ -765,6 +832,238 @@ class CallHandlers:
     # User-defined function calls
     # ------------------------------------------------------------------
 
+    def _func_local_length_defs(self, func_def) -> dict[str, str]:
+        """Collect a user function's local scalar length-vars to their RHS
+        expression string, e.g. ``qqeCalc`` with ``wp = sf * 2 - 1`` returns
+        ``{"wp": "sf * 2 - 1"}``.
+
+        Only plain (non-``var``/``varip``) declarations whose RHS is a pure
+        arithmetic expression over identifiers/numbers (NumberLiteral, Identifier,
+        BinOp, UnaryOp, or a math.* FuncCall) qualify — these are the shapes that
+        can legitimately feed a TA constructor length. Series-valued locals (whose
+        RHS is a ta.* call, a subscript, a ternary, etc.) are skipped so we never
+        inline a price series into a ctor-length slot. Names reassigned with ``:=``
+        are also skipped (their value is not a stable compile-time length).
+        """
+        def _is_arith(n) -> bool:
+            if isinstance(n, (NumberLiteral, Identifier)):
+                return True
+            if isinstance(n, BinOp):
+                return _is_arith(n.left) and _is_arith(n.right)
+            if isinstance(n, UnaryOp):
+                return _is_arith(n.operand)
+            if isinstance(n, Ternary):
+                # ``cond ? a : b`` — expand when both branches are arith.
+                # The condition may be a comparison/logical over arith leaves;
+                # rely on the codegen stability gate to reject series deps.
+                return (_is_arith(n.true_val) and _is_arith(n.false_val)
+                        and _is_arith(n.condition))
+            if isinstance(n, MemberAccess):
+                # ``timeframe.*`` / ``syminfo.*`` / ``math.pi`` etc. — stable
+                # per-run scalars that may appear inside a function-local
+                # derived length. Let the codegen stability classifier decide.
+                return True
+            if isinstance(n, FuncCall):
+                callee = n.callee
+                # Allow math.* helpers (math.round/sqrt/cos/...) over arith args.
+                if (isinstance(callee, MemberAccess)
+                        and isinstance(callee.object, Identifier)
+                        and callee.object.name == "math"):
+                    return all(_is_arith(a) for a in n.args)
+                # Pine type-cast builtins ``int(x)`` / ``float(x)`` / ``bool(x)``
+                # / ``string(x)`` — transparent over arith args. Common in
+                # derived TA lengths (``int(math.round(2 / a))``).
+                if (isinstance(callee, Identifier)
+                        and callee.name in ("int", "float", "bool", "string")):
+                    return all(_is_arith(a) for a in n.args)
+            return False
+
+        reassigned: set[str] = set()
+        def _scan_reassign(stmts):
+            from ..ast_nodes import Assignment
+            for s in stmts or []:
+                if isinstance(s, Assignment) and isinstance(s.target, Identifier):
+                    reassigned.add(s.target.name)
+                for attr in ("body", "else_body"):
+                    sub = getattr(s, attr, None)
+                    if isinstance(sub, list):
+                        _scan_reassign(sub)
+        _scan_reassign(func_def.body)
+
+        defs: dict[str, str] = {}
+        for stmt in func_def.body or []:
+            if (isinstance(stmt, VarDecl)
+                    and not stmt.is_var and not stmt.is_varip
+                    and stmt.name not in reassigned
+                    and stmt.value is not None
+                    and _is_arith(stmt.value)):
+                defs[stmt.name] = self._expr_to_str(stmt.value)
+        return defs
+
+    def _materialize_user_func_call_site_state(
+            self, func_name: str, cs_idx: int, node: FuncCall,
+            *, reuse_existing_owner: str | None = None) -> None:
+        """Materialize TA/fixnan state for one UDF call-site variant.
+
+        Ordinary call sites are handled while walking the AST.  A second class
+        is discovered only after that walk: a stateful helper reached through a
+        multi-call-site parent needs the parent's additional call-path indices
+        even though the helper has only one textual call.  The late propagation
+        pass in ``Analyzer._propagate_call_site_counts`` calls this same helper
+        for those inherited variants so the exported count never references a
+        TA/fixnan clone that was not actually declared.
+
+        ``reuse_existing_owner`` is used only by late propagation.  A
+        range-widened parent may already have materialized the default
+        ``{member}_cs{idx}`` clone for the borrowed callee site; when that clone
+        belongs to the parent currently being propagated, it is the desired
+        call-path state and must be reused rather than duplicated under a
+        disambiguated-but-unused name.
+        """
+        func_def = self._func_defs[func_name]
+
+        param_arg_map: dict[str, str] = {}
+        for p_idx, param_name in enumerate(func_def.params):
+            if p_idx < len(node.args):
+                param_arg_map[param_name] = self._expr_to_str(node.args[p_idx])
+
+        if func_name in self._func_ta_ranges:
+            start, end = self._func_ta_ranges[func_name]
+
+            # Map local derived length variables back to expressions over the
+            # function's parameters before substituting call-site arguments.
+            local_defs = self._func_local_length_defs(func_def)
+
+            def _subst_params(arg: str, pmap: dict[str, str]) -> str:
+                import re
+                result = arg
+                for param, value in sorted(
+                        pmap.items(), key=lambda item: len(item[0]), reverse=True):
+                    result = re.sub(rf'\b{re.escape(param)}\b', value, result)
+                return result
+
+            def _expand_locals(arg: str) -> str:
+                import re
+                if not local_defs:
+                    return arg
+                for _ in range(32):
+                    def _rep(match: re.Match) -> str:
+                        name = match.group(0)
+                        if name in local_defs:
+                            return "(" + local_defs[name] + ")"
+                        return name
+                    expanded = re.sub(r"[A-Za-z_][A-Za-z_0-9]*", _rep, arg)
+                    if expanded == arg:
+                        break
+                    arg = expanded
+                return arg
+
+            import re as _re
+            enclosing_params: set[str] = set()
+            for names in self._enclosing_func_params:
+                enclosing_params |= names
+
+            if cs_idx == 0:
+                # cs0 owns the source-level sites. Preserve their parameterized
+                # ctor args for every later direct or inherited clone.
+                for i in range(start, end):
+                    site = self._ta_call_sites[i]
+                    if not hasattr(site, '_orig_ctor_args'):
+                        site._orig_ctor_args = [
+                            _expand_locals(arg) for arg in site.ctor_args
+                        ]
+                    site.ctor_args = [
+                        _subst_params(arg, param_arg_map)
+                        for arg in site._orig_ctor_args
+                    ]
+                    # If a ctor is now expressed in an enclosing UDF's params,
+                    # retain that expression so the enclosing call can resolve
+                    # it and widen the enclosing TA range as before.
+                    if enclosing_params and self._nested_ta_touched is not None:
+                        for arg in site.ctor_args:
+                            tokens = set(_re.findall(
+                                r"[A-Za-z_][A-Za-z_0-9]*", arg))
+                            if tokens & enclosing_params:
+                                site._orig_ctor_args = list(site.ctor_args)
+                                self._nested_ta_touched.add(i)
+                                break
+            else:
+                clone_name_map: dict[str, str] = {}
+                for i in range(start, end):
+                    orig = self._ta_call_sites[i]
+                    orig_args = getattr(orig, '_orig_ctor_args', orig.ctor_args)
+                    resolved_ctor = [
+                        _subst_params(arg, param_arg_map) for arg in orig_args
+                    ]
+                    clone_name = f"{orig.member_name}_cs{cs_idx}"
+                    existing = next(
+                        (site for site in self._ta_call_sites
+                         if site.member_name == clone_name),
+                        None,
+                    )
+                    if (reuse_existing_owner is not None
+                            and existing is not None
+                            and existing.owner_func == reuse_existing_owner):
+                        # The active parent's widened range already made the
+                        # exact member this inherited callee variant needs.
+                        continue
+                    if clone_name in self._ta_member_names:
+                        base = clone_name
+                        suffix = 2
+                        while clone_name in self._ta_member_names:
+                            clone_name = f"{base}_u{suffix}"
+                            suffix += 1
+                        clone_name_map[orig.member_name] = clone_name
+                    cloned = TACallSite(
+                        member_name=clone_name,
+                        class_name=orig.class_name,
+                        ctor_args=resolved_ctor,
+                        compute_args=orig.compute_args[:],
+                        returns_tuple=orig.returns_tuple,
+                        node=orig.node,
+                        is_static=orig.is_static,
+                        owner_func=func_name,
+                    )
+                    self._ta_call_sites.append(cloned)
+                    self._ta_member_names.add(clone_name)
+                if clone_name_map:
+                    self._func_cs_ta_clone_names[(func_name, cs_idx)] = clone_name_map
+
+        # fixnan is stateful for the same reason as a rolling TA reducer: each
+        # emitted function variant needs its own previous-value member.
+        fn_indices = self._func_fixnan_indices.get(func_name, [])
+        if cs_idx > 0 and fn_indices:
+            clone_map: dict[str, str] = {}
+            for fi in fn_indices:
+                orig = self._fixnan_sites[fi]
+                clone_name = f"{orig.member_name}_cs{cs_idx}"
+                existing = next(
+                    (site for site in self._fixnan_sites
+                     if site.member_name == clone_name),
+                    None,
+                )
+                if (reuse_existing_owner is not None
+                        and existing is not None
+                        and existing.owner_func == reuse_existing_owner):
+                    continue
+                if clone_name in self._fixnan_member_names:
+                    base = clone_name
+                    suffix = 2
+                    while clone_name in self._fixnan_member_names:
+                        clone_name = f"{base}_u{suffix}"
+                        suffix += 1
+                    clone_map[orig.member_name] = clone_name
+                cloned = FixnanCallSite(
+                    member_name=clone_name,
+                    pine_type=orig.pine_type,
+                    node=orig.node,
+                    owner_func=func_name,
+                )
+                self._fixnan_sites.append(cloned)
+                self._fixnan_member_names.add(clone_name)
+            if clone_map:
+                self._func_cs_fixnan_clone_names[(func_name, cs_idx)] = clone_map
+
     def _handle_user_func_call(self, func_name: str, node: FuncCall) -> PineType:
         """Handle calls to user-defined functions."""
         func_def = self._func_defs[func_name]
@@ -783,14 +1082,28 @@ class CallHandlers:
         # For now, use the cached return type from initial analysis
         return_type = self._func_return_types.get(func_name, PineType.FLOAT)
 
-        # If the return type was UNKNOWN or VOID, infer from param types
+        # If the return type was UNKNOWN or VOID, infer it ONLY when the body
+        # is a single bare identifier that returns a parameter directly
+        # (``f(s) => s``). Inferring from params for arbitrary bodies misfires
+        # when a function merely HAS a string param but returns something else
+        # (e.g. ``getLineStyle(s) => switch s ... => line.style_solid`` or a
+        # body ending in ``label.new(...)``). Other cases rely on the cached
+        # body type plus udt_return_type / tuple inference.
         if return_type in (PineType.UNKNOWN, PineType.VOID):
-            if any(t == PineType.STRING for t in param_types):
-                return_type = PineType.STRING
-            elif any(t == PineType.FLOAT for t in param_types):
-                return_type = PineType.FLOAT
-            elif any(t == PineType.INT for t in param_types):
-                return_type = PineType.INT
+            if (func_def.is_single_expr and func_def.body
+                    and isinstance(func_def.body[0], ExprStmt)
+                    and isinstance(func_def.body[0].expr, Identifier)):
+                ret_name = func_def.body[0].expr.name
+                for idx, pname in enumerate(func_def.params):
+                    if pname == ret_name and idx < len(param_types):
+                        pt = param_types[idx]
+                        if pt == PineType.STRING:
+                            return_type = PineType.STRING
+                        elif pt == PineType.INT:
+                            return_type = PineType.INT
+                        elif pt == PineType.FLOAT:
+                            return_type = PineType.FLOAT
+                        break
 
         # If this function has series params, ensure bar-field arguments
         # passed at the call site are registered as series_bar_fields so that
@@ -802,66 +1115,29 @@ class CallHandlers:
                     arg = node.args[p_idx]
                     if isinstance(arg, Identifier) and arg.name in BAR_FIELDS:
                         self._series_bar_fields.add(arg.name)
-
-        # Per-call-site cloning: if this function has TA calls or series vars,
-        # track call sites so codegen can create per-call-site variants.
-        # This prevents shared state corruption when the function is called
-        # multiple times per bar.
+                    elif isinstance(arg, Identifier):
+                        sym = self._symbols.resolve(arg.name)
+                        spec = getattr(sym, "type_spec", None) if sym is not None else None
+                        if spec is not None and spec.kind in ("array", "map", "matrix"):
+                            continue
+                        if sym is not None:
+                            sym.is_series = True
+                            if sym.scope and sym.scope.startswith("func_"):
+                                caller_name = sym.scope[5:]
+                                self._func_series_vars.setdefault(caller_name, set()).add(arg.name)
+                            else:
+                                self._series_vars.add(arg.name)
+
+        # Per-call-site cloning: TA, series/var, and fixnan state all advance
+        # across bars/calls and therefore require isolated UDF variants.
         has_ta = func_name in self._func_ta_ranges
         has_series = func_name in self._func_series_vars or func_name in self._func_var_members
-        if has_ta or has_series:
+        has_fixnan = func_name in self._func_fixnan_indices
+        if has_ta or has_series or has_fixnan:
             cs_idx = self._func_call_site_count.get(func_name, 0)
             self._func_call_site_count[func_name] = cs_idx + 1
             self._func_call_cs_map[id(node)] = (func_name, cs_idx)
-
-            # Build parameter -> call-site argument string mapping
-            param_arg_map: dict[str, str] = {}
-            for p_idx, param_name in enumerate(func_def.params):
-                if p_idx < len(node.args):
-                    param_arg_map[param_name] = self._expr_to_str(node.args[p_idx])
-
-            # Clone TA call sites (only if function has TA ranges)
-            if has_ta:
-                start, end = self._func_ta_ranges[func_name]
-
-                def _subst_params(arg: str, pmap: dict[str, str]) -> str:
-                    """Substitute parameter names in an expression string.
-
-                    Handles both exact matches ('len' -> 'len3') and parameter
-                    names within expressions ('len / 2' -> 'len3 / 2').
-                    """
-                    import re
-                    result = arg
-                    # Sort by length descending to avoid partial replacements
-                    for param, value in sorted(pmap.items(), key=lambda x: len(x[0]), reverse=True):
-                        result = re.sub(rf'\b{re.escape(param)}\b', value, result)
-                    return result
-
-                if cs_idx == 0:
-                    # First call site: save original param-based ctor_args for future cloning,
-                    # then resolve to actual call-site values
-                    for i in range(start, end):
-                        site = self._ta_call_sites[i]
-                        if not hasattr(site, '_orig_ctor_args'):
-                            site._orig_ctor_args = site.ctor_args[:]
-                        site.ctor_args = [_subst_params(a, param_arg_map) for a in site._orig_ctor_args]
-                else:
-                    # Subsequent call sites: clone using saved original param names,
-                    # substituted with this call site's arguments
-                    for i in range(start, end):
-                        orig = self._ta_call_sites[i]
-                        orig_args = getattr(orig, '_orig_ctor_args', orig.ctor_args)
-                        resolved_ctor = [_subst_params(a, param_arg_map) for a in orig_args]
-                        cloned = TACallSite(
-                            member_name=f"{orig.member_name}_cs{cs_idx}",
-                            class_name=orig.class_name,
-                            ctor_args=resolved_ctor,
-                            compute_args=orig.compute_args[:],
-                            returns_tuple=orig.returns_tuple,
-                            node=orig.node,
-                            is_static=orig.is_static,
-                        )
-                        self._ta_call_sites.append(cloned)
+            self._materialize_user_func_call_site_state(func_name, cs_idx, node)
 
         # Create or update FuncInfo
         is_tuple = self._func_returns_tuple.get(func_name, False)
@@ -869,6 +1145,18 @@ class CallHandlers:
         # Forward UDT-return inference (set in _visit_FuncDef) so codegen can
         # emit the struct return type. Probe: udt-method-probe-20.
         udt_ret = self._func_udt_return_types.get(func_name)
+        ret_spec = getattr(self, "_func_return_type_specs", {}).get(func_name)
+        # Per-param TypeSpec: declared hints are authoritative; for untyped
+        # params, infer from the call-site argument's type_spec (so an untyped
+        # ``s`` used as a string, or a UDT passed by value, emits correctly).
+        param_specs = list(
+            getattr(self, "_func_param_type_specs", {}).get(func_name)
+            or self._param_type_specs_from_def(func_def)
+        )
+        arg_specs = [self._type_spec_from_expr(arg) for arg in node.args]
+        for i in range(len(param_specs)):
+            if param_specs[i] is None and i < len(arg_specs):
+                param_specs[i] = arg_specs[i]
         existing = [fi for fi in self._func_infos if fi.name == func_name]
         if not existing:
             fi = FuncInfo(
@@ -879,6 +1167,8 @@ class CallHandlers:
                 returns_tuple=is_tuple,
                 tuple_element_count=tuple_count,
                 udt_return_type=udt_ret,
+                param_type_specs=param_specs,
+                return_type_spec=ret_spec,
             )
             self._func_infos.append(fi)
         else:
@@ -889,7 +1179,17 @@ class CallHandlers:
             for i, pt in enumerate(param_types):
                 if i < len(fi.param_types) and fi.param_types[i] == PineType.UNKNOWN:
                     fi.param_types[i] = pt
+            # Merge per-param TypeSpecs: keep declared hints (authoritative),
+            # fill untyped slots from this call site if still unknown.
+            if not fi.param_type_specs:
+                fi.param_type_specs = list(param_specs)
+            else:
+                for i in range(len(param_specs)):
+                    if i < len(fi.param_type_specs) and fi.param_type_specs[i] is None:
+                        fi.param_type_specs[i] = param_specs[i]
             if fi.udt_return_type is None and udt_ret is not None:
                 fi.udt_return_type = udt_ret
+            if fi.return_type_spec is None and ret_spec is not None:
+                fi.return_type_spec = ret_spec
 
         return return_type