@hoox-sh/pynets 0.2.0 → 0.3.0

package/src/runtime/compile/emit.ts

@@ -8,6 +8,7 @@ import { parse } from "../../ast/helper.ts";
 import type {
   Assign,
   Attribute,
+  AugAssign,
   BinOp,
   BoolOp,
   Call,
@@ -23,6 +24,7 @@ import type {
   If,
   Import,
   Name,
+  Once,
   ReAssign,
   Script,
   stmt,
@@ -33,7 +35,7 @@ import type {
   UnaryOp,
   While,
 } from "../../ast/nodes.ts";
-import { emitCall } from "./emit_call.ts";
+import { emitCall, TA_BARE_ATTRS } from "./emit_call.ts";
 import { emitFunctionDef } from "./emit_udf.ts";
 import type { EmitCtx, VisitFn } from "./types.ts";
 
@@ -47,6 +49,7 @@ const CHART_LOCALS = new Set([
   "close",
   "volume",
   "time",
+  "time_close",
   "bar_index",
   "last_bar_index",
   "na",
@@ -138,6 +141,7 @@ export function newEmitCtx(): EmitCtx {
     usesStrategy: false,
     nextSite: 0,
     errors: [],
+    onceFlags: [],
     currentFunc: null,
     currentParamNames: new Set(),
     currentSeriesParams: new Set(),
@@ -210,6 +214,9 @@ function assemble(ctx: EmitCtx, body: string[]): string {
     lines.push(`  const ${arr} = new Array(n_bars).fill(null);`);
   }
   lines.push("  const __var_inited = Object.create(null);");
+  for (const flag of ctx.onceFlags) {
+    lines.push(`  let ${flag} = false;`);
+  }
   for (const fn of ctx.functions) {
     if (fn) lines.push(indentBlock(fn, 2));
   }
@@ -223,6 +230,7 @@ function assemble(ctx: EmitCtx, body: string[]): string {
   lines.push("    const close = close_arr[__bar_idx];");
   lines.push("    const volume = vol_arr[__bar_idx];");
   lines.push("    const time = time_arr[__bar_idx];");
+  lines.push("    const time_close = __h.timeClose(__bar_idx, time_arr);");
   lines.push("    const bar_index = __bar_idx;");
   lines.push("    const last_bar_index = n_bars - 1;");
   lines.push("    const na = null;");
@@ -267,6 +275,7 @@ function emitNode(state: State, node: expr | stmt | null | undefined): string {
   switch (node.kind) {
     case "Assign":
     case "ReAssign":
+    case "AugAssign":
     case "FunctionDef":
     case "TypeDef":
     case "EnumDef":
@@ -274,6 +283,7 @@ function emitNode(state: State, node: expr | stmt | null | undefined): string {
     case "Break":
     case "Continue":
     case "Expr":
+    case "Once":
       return emitStmt(state, node);
     default:
       return emitExpr(state, node);
@@ -288,6 +298,8 @@ function emitStmt(state: State, node: stmt | expr): string {
       return emitAssign(state, node);
     case "ReAssign":
       return emitReAssign(state, node);
+    case "AugAssign":
+      return emitAugAssign(state, node);
     case "FunctionDef":
       emitFunctionDef(state.ctx, node as FunctionDef, (s) => emitStmt(state, s));
       if (state.ctx.currentFunc != null) state.ctx.currentFunc = null;
@@ -312,6 +324,8 @@ function emitStmt(state: State, node: stmt | expr): string {
       return emitWhile(state, node);
     case "Switch":
       return emitSwitch(state, node, false);
+    case "Once":
+      return emitOnce(state, node);
     default:
       return emitExprStmt(state, node);
   }
@@ -330,6 +344,8 @@ function emitExprStmt(state: State, value: expr | null | undefined): string {
       return emitWhile(state, value);
     case "Switch":
       return emitSwitch(state, value, false);
+    case "Once":
+      return emitOnce(state, value);
     default: {
       const e = emitExpr(state, value);
       if (!e || e === "null") return "";
@@ -373,6 +389,12 @@ function emitExpr(state: State, node: expr | null | undefined): string {
       return iife(emitForIn(state, node), "null");
     case "While":
       return iife(emitWhile(state, node), "null");
+    case "Qualify":
+      return emitExpr(state, node.value);
+    case "Specialize":
+      return emitExpr(state, node.value);
+    case "AugAssign":
+      return emitAugAssign(state, node);
     default:
       return "null";
   }
@@ -423,11 +445,25 @@ function emitBinOp(state: State, node: BinOp): string {
       return `__h.div(${left}, ${right})`;
     case "Mod":
       return `__h.mod(${left}, ${right})`;
+    case "BitAnd":
+      return emitBitOp(left, right, "&");
+    case "BitOr":
+      return emitBitOp(left, right, "|");
+    case "BitXor":
+      return emitBitOp(left, right, "^");
+    case "LShift":
+      return emitBitOp(left, right, "<<");
+    case "RShift":
+      return emitBitOp(left, right, ">>");
     default:
       return `__h.add(${left}, ${right})`;
   }
 }
 
+function emitBitOp(left: string, right: string, op: "&" | "|" | "^" | "<<" | ">>"): string {
+  return `((__l, __r) => { if (__h.isNa(__l) || __h.isNa(__r)) return null; return (__h.nz(__l) | 0) ${op} (__h.nz(__r) | 0); })(${left}, ${right})`;
+}
+
 function emitUnaryOp(state: State, node: UnaryOp): string {
   const operand = emitExpr(state, node.operand);
   switch (opKind(node.op)) {
@@ -500,9 +536,13 @@ const STRATEGY_QUERY_ATTRS = new Set([
   "openprofit",
   "opentrades",
   "closedtrades",
+  "leverage",
+  "margin_liquidation_price",
+  "position_avg_price",
+  "initial_capital",
 ]);
 
-/** Namespaces used only as call receivers — not UDT field access. */
+/** Namespaces used as call receivers or dotted constants — not UDT field access. */
 const ATTR_NAMESPACES = new Set([
   "ta",
   "math",
@@ -521,8 +561,150 @@ const ATTR_NAMESPACES = new Set([
   "session",
   "chart",
   "log",
+  "order",
+  "format",
+  "text",
+  "size",
+  "barmerge",
+  "shape",
+  "location",
+  "xloc",
+  "yloc",
+  "extend",
+  "display",
+  "position",
+  "hline",
+  "dayofweek",
+  "month",
+]);
+
+/** Dotted keys from Python `_MATH_CONSTANTS` (base.py). Unknown member → na. */
+const ATTR_CONSTANTS: Record<string, string | number | boolean> = {
+  // Math constants — Python SoT `_MATH_CONSTANTS` (evaluator/base.py).
+  "math.pi": Math.PI,
+  "math.e": Math.E,
+  "math.phi": (1 + Math.sqrt(5)) / 2,
+  "math.rphi": 2 / (1 + Math.sqrt(5)),
+  "format.mintick": "mintick",
+  "format.percent": "percent",
+  "format.volume": "volume",
+  "format.price": "price",
+  "text.formatting.none": "",
+  "text.formatting.bold": "bold",
+  "text.formatting.italic": "italic",
+  "text.formatting.bold_italic": "bold italic",
+  "size.auto": "auto",
+  "size.tiny": 8,
+  "size.small": 10,
+  "size.normal": 12,
+  "size.large": 16,
+  "size.huge": 20,
+  "order.ascending": 1,
+  "order.descending": -1,
+  "barmerge.gaps_on": true,
+  "barmerge.gaps_off": false,
+  "barmerge.lookahead_on": true,
+  "barmerge.lookahead_off": false,
+  "shape.arrowup": "arrowup",
+  "shape.arrowdown": "arrowdown",
+  "shape.circle": "circle",
+  "shape.cross": "cross",
+  "shape.diamond": "diamond",
+  "shape.flag": "flag",
+  "shape.labelup": "labelup",
+  "shape.labeldown": "labeldown",
+  "shape.square": "square",
+  "shape.triangledown": "triangledown",
+  "shape.triangleup": "triangleup",
+  "shape.xcross": "xcross",
+  "location.abovebar": "abovebar",
+  "location.belowbar": "belowbar",
+  "location.top": "top",
+  "location.bottom": "bottom",
+  "location.absolute": "absolute",
+  "xloc.bar_index": "bar_index",
+  "xloc.bar_time": "bar_time",
+  "yloc.price": "price",
+  "yloc.abovebar": "abovebar",
+  "yloc.belowbar": "belowbar",
+  "extend.none": "none",
+  "extend.left": "left",
+  "extend.right": "right",
+  "extend.both": "both",
+  // Pine display.* is a bitfield: none=0 pane=1 data_window=2 price_scale=4
+  // status_line=8 all=15. Integers so `display.pane + display.data_window` works
+  // (Python SoT base.py / compiler.py — matches interpret backend).
+  "display.none": 0,
+  "display.pane": 1,
+  "display.all": 15,
+  "display.data_window": 2,
+  "display.price_scale": 4,
+  "display.status_line": 8,
+  "position.top_left": "top_left",
+  "position.top_center": "top_center",
+  "position.top_right": "top_right",
+  "position.middle_left": "middle_left",
+  "position.middle_center": "middle_center",
+  "position.middle_right": "middle_right",
+  "position.bottom_left": "bottom_left",
+  "position.bottom_center": "bottom_center",
+  "position.bottom_right": "bottom_right",
+  "hline.style_solid": "solid",
+  "hline.style_dashed": "dashed",
+  "hline.style_dotted": "dotted",
+  "dayofweek.sunday": 1,
+  "dayofweek.monday": 2,
+  "dayofweek.tuesday": 3,
+  "dayofweek.wednesday": 4,
+  "dayofweek.thursday": 5,
+  "dayofweek.friday": 6,
+  "dayofweek.saturday": 7,
+  "month.january": 1,
+  "month.february": 2,
+  "month.march": 3,
+  "month.april": 4,
+  "month.may": 5,
+  "month.june": 6,
+  "month.july": 7,
+  "month.august": 8,
+  "month.september": 9,
+  "month.october": 10,
+  "month.november": 11,
+  "month.december": 12,
+};
+
+const ATTR_CONST_NS = new Set([
+  "math",
+  "order",
+  "format",
+  "text",
+  "size",
+  "barmerge",
+  "shape",
+  "location",
+  "xloc",
+  "yloc",
+  "extend",
+  "display",
+  "position",
+  "hline",
+  "dayofweek",
+  "month",
 ]);
 
+function attrQualifiedName(node: Attribute): string | null {
+  const parts: string[] = [node.attr];
+  let cur: expr = node.value;
+  while (cur.kind === "Attribute") {
+    parts.push(cur.attr);
+    cur = cur.value;
+  }
+  if (cur.kind !== "Name") return null;
+  parts.push(cur.id);
+  parts.reverse();
+  return parts.join(".");
+}
+
 function emitBarstate(attr: string): string {
   switch (attr) {
     case "isfirst":
@@ -584,19 +766,34 @@ function emitChart(attr: string): string {
   }
 }
 
+// Static defaults mirror the Python SoT flat table in
+// pynescript/ast/evaluator/base.py ("Chart timeframe defaults (daily)"),
+// cross-checked with runtime/host.py Timeframe (same values). The compile host
+// has no timeframe context (`__h.timeframe` is undefined), so these are the
+// effective values; interpret derives them dynamically from env.timeframe and
+// resolves to the same daily table when no timeframe is configured — every
+// timeframe.* member is backend-parity (commit 6977a42; pinned in
+// test/compile_timeframe_defaults.test.ts).
 function emitTimeframe(attr: string): string {
   switch (attr) {
     case "period":
-      return `(__h.timeframe || "1")`;
+    case "main_period":
+      return `(__h.timeframe || "D")`;
     case "multiplier":
       return "1";
     case "isintraday":
-      return "true";
-    case "isdaily":
     case "isweekly":
     case "ismonthly":
+    case "isseconds":
+    case "isinseconds":
+    case "isminutes":
+    case "ishours":
       return "false";
+    case "isdaily":
+    case "isdwm":
+      return "true";
     default:
+      // Unknown member → na (documented behavior).
       return "null";
   }
 }
@@ -607,6 +804,8 @@ function emitSyminfo(attr: string): string {
     case "tickerid":
     case "root":
       return `"SYMBOL"`;
+    case "prefix":
+      return `__h.syminfoPrefix()`;
     default:
       return "null";
   }
@@ -635,10 +834,22 @@ function emitAttribute(state: State, node: Attribute): string {
   if (node.value.kind === "Name" && node.value.id === "color") {
     return `__h.colorByName(${JSON.stringify(node.attr)})`;
   }
+  if (node.value.kind === "Name" && node.value.id === "ta") {
+    // Bare `ta.<attr>` auto-calls only for the whitelist interpret's
+    // evalTaAttr supports; everything else is na on both backends.
+    if (!TA_BARE_ATTRS.has(node.attr)) return "null";
+    return emitCall(state.ctx, { kind: "Call", func: node, args: [] }, visitOf(state));
+  }
   if (node.value.kind === "Name" && node.value.id === "strategy") {
     const attr = node.attr;
     if (attr === "long") return `"long"`;
     if (attr === "short") return `"short"`;
+    if (attr === "avg_price_stock") return `"stock"`;
+    if (attr === "avg_price_futures") return `"futures"`;
+    if (attr === "avg_price_inverse") return `"inverse"`;
+    if (attr === "cash") return `"cash"`;
+    if (attr === "fixed") return `"fixed"`;
+    if (attr === "percent_of_equity") return `"percent_of_equity"`;
     if (STRATEGY_QUERY_ATTRS.has(attr)) {
       state.ctx.usesStrategy = true;
       return `__h.strategy.${attr}()`;
@@ -650,6 +861,15 @@ function emitAttribute(state: State, node: Attribute): string {
       return JSON.stringify(node.attr);
     }
   }
+  const qn = attrQualifiedName(node);
+  if (qn != null) {
+    const dot = qn.indexOf(".");
+    const root = dot < 0 ? qn : qn.slice(0, dot);
+    if (ATTR_CONST_NS.has(root)) {
+      const lit = ATTR_CONSTANTS[qn];
+      return lit === undefined ? "null" : JSON.stringify(lit);
+    }
+  }
   const obj = emitExpr(state, node.value);
   if (node.value.kind === "Name" && state.ctx.udtTypes.has(node.value.id)) {
     return `__h.udtGet(${obj}, ${JSON.stringify(node.attr)})`;
@@ -772,6 +992,21 @@ function emitReAssign(state: State, node: ReAssign): string {
   return emitStore(state, node.target, emitExpr(state, node.value), false);
 }
 
+function emitAugAssign(state: State, node: AugAssign): string {
+  // Attribute targets store the raw rhs — the old field is discarded
+  // (quirk preserved from Python visit_AugAssign and interpret evalAugAssign).
+  if (node.target.kind === "Attribute") {
+    return emitStore(state, node.target, emitExpr(state, node.value), false);
+  }
+  const rhs = emitBinOp(state, {
+    kind: "BinOp",
+    left: node.target,
+    op: node.op,
+    right: node.value,
+  });
+  return emitStore(state, node.target, rhs, false);
+}
+
 function emitStore(state: State, target: expr, rhs: string, isVar: boolean): string {
   if (target.kind === "Name") {
     return storeName(state, target.id, rhs, isVar);
@@ -842,6 +1077,15 @@ function emitIfStmt(state: State, node: If, asExpr = false): string {
   return `if (${test}) {\n${indentBlock(thenBlock, 2)}\n} else {\n${indentBlock(elseBlock, 2)}\n}`;
 }
 
+function emitOnce(state: State, node: Once): string {
+  const flag = `__once_${state.ctx.onceFlags.length}`;
+  state.ctx.onceFlags.push(flag);
+  const test = node.test == null ? "true" : `__h.nz(${emitExpr(state, node.test)})`;
+  const body = emitBranch(state, node.body, false);
+  const inner = body ? `${indentBlock(body, 4)}\n    ${flag} = true;` : `    ${flag} = true;`;
+  return `if (!${flag}) {\n  if (${test}) {\n${inner}\n  }\n}`;
+}
+
 function emitBranch(state: State, stmts: stmt[], asExpr: boolean): string {
   const list = stmts ?? [];
   if (list.length === 0) return asExpr ? "return null;" : "";
@@ -900,6 +1144,7 @@ function isValueExpr(node: expr): boolean {
     case "ForIn":
     case "While":
     case "Switch":
+    case "Once":
       return false;
     default:
       return true;
@@ -1021,10 +1266,7 @@ function emitSwitch(state: State, node: Switch, asExpr: boolean): string {
       defaultCase = c;
       continue;
     }
-    const cond =
-      subjN != null
-        ? `__h.eq(${subjN}, ${emitExpr(state, c.pattern)})`
-        : `__h.nz(${emitExpr(state, c.pattern)})`;
+    const cond = emitSwitchCond(state, c.pattern, subjN);
     const kw = first ? "if" : "else if";
     first = false;
     const body = emitBranch(state, c.body ?? [], asExpr);
@@ -1045,6 +1287,21 @@ function emitSwitch(state: State, node: Switch, asExpr: boolean): string {
   return `{\n${indentBlock(block, 2)}\n}`;
 }
 
+/** Python `_switch_case_matches`: a comma arm (`1, 2 =>`) matches any element. */
+function emitSwitchCond(state: State, pattern: expr, subjN: string | null): string {
+  if (pattern.kind === "Tuple") {
+    const parts = (pattern.elts ?? []).map((e) => emitSwitchLeaf(state, e, subjN));
+    return parts.length ? `(${parts.join(" || ")})` : "false";
+  }
+  return emitSwitchLeaf(state, pattern, subjN);
+}
+
+function emitSwitchLeaf(state: State, pattern: expr, subjN: string | null): string {
+  return subjN != null
+    ? `__h.eq(${subjN}, ${emitExpr(state, pattern)})`
+    : `__h.nz(${emitExpr(state, pattern)})`;
+}
+
 function emitBlock(state: State, stmts: stmt[] | null | undefined): string {
   const lines: string[] = [];
   for (const s of stmts ?? []) {