From ce54e817b799827ba75c5e761c910a0a75494ef9 Mon Sep 17 00:00:00 2001 From: Martin Vogel Date: Sat, 9 May 2026 09:53:43 +0200 Subject: [PATCH] feat(py_lsp): subscript + super().__init__() + operator dunders MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round 5 parity push: Subscript value typing — py_eval_expr_type's subscript branch resolves indexing against the container's TEMPLATE type: - TEMPLATE("dict"|"Mapping"|"defaultdict"|"OrderedDict"|...)[K] -> V - TEMPLATE("list"|"set"|"frozenset"|"Sequence"|"Iterable"|...)[i] -> elem - TEMPLATE("tuple", [A, B, ...])[literal_int] -> args[idx] - Non-literal index returns UNION of all element types. This unlocks the `self.cache[k].method()` pattern that's everywhere in real Python. Also fixes `items[0].method()` for typed lists. super().__init__(...) constructor edge — super(). already worked for regular methods; __init__ is special-cased because parent classes typically don't register __init__ with a useful return type. Strategy "lsp_super_init" with confidence 0.85. Operator dunders — a + b looks up __add__ on a's type and uses its return type. Same for *, -, /, //, %, **, <<, >>, &, |, ^, @. Falls back to left-operand type when registry has no info (correct for the common int+int / str+str / list+list cases). 3 new tests; all 2892 prior tests stay green. Bench unchanged at 95% (fixture doesn't exercise these patterns directly — Round 6 will expand it). --- internal/cbm/lsp/py_lsp.c | 86 +++++++++++++++++++++++++++++++++++++++ tests/test_py_lsp.c | 56 +++++++++++++++++++++++++ 2 files changed, 142 insertions(+) diff --git a/internal/cbm/lsp/py_lsp.c b/internal/cbm/lsp/py_lsp.c index e6a2303a1..97929be65 100644 --- a/internal/cbm/lsp/py_lsp.c +++ b/internal/cbm/lsp/py_lsp.c @@ -603,7 +603,38 @@ static const CBMType* py_eval_expr_type(PyLSPContext* ctx, TSNode node) { if (strcmp(k, "binary_operator") == 0) { TSNode left = ts_node_child_by_field_name(node, "left", 4); + TSNode op = ts_node_child_by_field_name(node, "operator", 8); const CBMType* lt = py_eval_expr_type(ctx, left); + // Try the operator's dunder method on left's type. Falls back to + // left-operand type when registry has no info (most common: int+int, + // str+str, list+list — typing-by-left is correct in those cases). + if (lt && lt->kind == CBM_TYPE_NAMED && !ts_node_is_null(op)) { + char* op_text = py_node_text(ctx, op); + const char* dunder = NULL; + if (op_text) { + if (strcmp(op_text, "+") == 0) dunder = "__add__"; + else if (strcmp(op_text, "-") == 0) dunder = "__sub__"; + else if (strcmp(op_text, "*") == 0) dunder = "__mul__"; + else if (strcmp(op_text, "/") == 0) dunder = "__truediv__"; + else if (strcmp(op_text, "//") == 0) dunder = "__floordiv__"; + else if (strcmp(op_text, "%") == 0) dunder = "__mod__"; + else if (strcmp(op_text, "**") == 0) dunder = "__pow__"; + else if (strcmp(op_text, "<<") == 0) dunder = "__lshift__"; + else if (strcmp(op_text, ">>") == 0) dunder = "__rshift__"; + else if (strcmp(op_text, "&") == 0) dunder = "__and__"; + else if (strcmp(op_text, "|") == 0) dunder = "__or__"; + else if (strcmp(op_text, "^") == 0) dunder = "__xor__"; + else if (strcmp(op_text, "@") == 0) dunder = "__matmul__"; + } + if (dunder) { + const CBMRegisteredFunc* f = py_lookup_attribute(ctx, + lt->data.named.qualified_name, dunder); + if (f) { + return py_func_return_type_recv(ctx, f->qualified_name, + lt->data.named.qualified_name); + } + } + } return lt ? lt : cbm_type_unknown(); } @@ -639,6 +670,51 @@ static const CBMType* py_eval_expr_type(PyLSPContext* ctx, TSNode node) { } } + // Subscript: container[index]. For TEMPLATE("dict", [K, V])[K] -> V; + // TEMPLATE("list"|"set"|...)[int] -> elem; TEMPLATE("tuple", [A, B, + // ...])[int] -> A (union when index isn't a literal int). + if (strcmp(k, "subscript") == 0) { + TSNode value = ts_node_child_by_field_name(node, "value", 5); + if (ts_node_is_null(value)) return cbm_type_unknown(); + const CBMType* container = py_eval_expr_type(ctx, value); + if (!container || container->kind != CBM_TYPE_TEMPLATE) { + return cbm_type_unknown(); + } + const char* tname = container->data.template_type.template_name; + const CBMType** args = container->data.template_type.template_args; + int n = container->data.template_type.arg_count; + if (!tname || !args || n <= 0) return cbm_type_unknown(); + if (strcmp(tname, "dict") == 0 || strcmp(tname, "Mapping") == 0 || + strcmp(tname, "MutableMapping") == 0 || strcmp(tname, "defaultdict") == 0 || + strcmp(tname, "OrderedDict") == 0 || strcmp(tname, "ChainMap") == 0) { + // Value type is the second template arg. + if (n >= 2 && args[1]) return args[1]; + return cbm_type_unknown(); + } + if (strcmp(tname, "list") == 0 || strcmp(tname, "set") == 0 || + strcmp(tname, "frozenset") == 0 || strcmp(tname, "Sequence") == 0 || + strcmp(tname, "MutableSequence") == 0 || strcmp(tname, "deque") == 0 || + strcmp(tname, "Iterable") == 0 || strcmp(tname, "Iterator") == 0) { + return args[0]; + } + if (strcmp(tname, "tuple") == 0) { + // Try to index by literal integer subscript. + TSNode sub = ts_node_child_by_field_name(node, "subscript", 9); + if (!ts_node_is_null(sub) && strcmp(ts_node_type(sub), "integer") == 0) { + char* idx_text = py_node_text(ctx, sub); + if (idx_text) { + int idx = atoi(idx_text); + if (idx >= 0 && idx < n) return args[idx]; + } + } + // Non-literal index: union of all elements. + if (n == 1) return args[0]; + return cbm_type_union(ctx->arena, args, n); + } + // Other generics: best-effort first arg. + return args[0]; + } + return cbm_type_unknown(); } @@ -786,6 +862,16 @@ static void py_emit_call_for(PyLSPContext* ctx, TSNode call_node) { "lsp_super", 0.88f); return; } + // Special case: super().__init__ — most parent + // classes don't register __init__ with a return, + // but we still want to emit the constructor edge. + if (strcmp(attr_name, "__init__") == 0) { + const char* init_qn = cbm_arena_sprintf(ctx->arena, + "%s.__init__", enclosing->embedded_types[i]); + py_emit_resolved_call(ctx, init_qn, + "lsp_super_init", 0.85f); + return; + } } } } diff --git a/tests/test_py_lsp.c b/tests/test_py_lsp.c index e067c6ef8..3316714f3 100644 --- a/tests/test_py_lsp.c +++ b/tests/test_py_lsp.c @@ -1093,6 +1093,58 @@ TEST(pylsp_round4_instance_attribute_class_annotation) { PASS(); } +/* ── Round 5 — subscript, super().__init__, operator dunders ───── */ + +TEST(pylsp_round5_dict_subscript_value_type) { + /* self.cache: dict[str, Foo]; self.cache[k].method() resolves. */ + CBMFileResult *r = extract_py( + "class Foo:\n" + " def method(self):\n" + " return 1\n" + "class C:\n" + " def __init__(self):\n" + " self.cache: dict[str, Foo] = {}\n" + " def use(self, k):\n" + " return self.cache[k].method()\n"); + ASSERT_NOT_NULL(r); + ASSERT_GTE(require_resolved(r, "use", "method"), 0); + cbm_free_result(r); + PASS(); +} + +TEST(pylsp_round5_list_subscript_value_type) { + CBMFileResult *r = extract_py( + "class Foo:\n" + " def method(self):\n" + " return 1\n" + "def use(items: list[Foo]):\n" + " return items[0].method()\n"); + ASSERT_NOT_NULL(r); + ASSERT_GTE(require_resolved(r, "use", "method"), 0); + cbm_free_result(r); + PASS(); +} + +TEST(pylsp_round5_super_init) { + CBMFileResult *r = extract_py( + "class Base:\n" + " def __init__(self, root):\n" + " self.root = root\n" + "class Child(Base):\n" + " def __init__(self, root, extra):\n" + " super().__init__(root)\n" + " self.extra = extra\n"); + ASSERT_NOT_NULL(r); + int idx = find_resolved(r, "__init__", "__init__"); + ASSERT_GTE(idx, 0); + if (idx >= 0) { + const CBMResolvedCall *rc = &r->resolved_calls.items[idx]; + ASSERT(strstr(rc->callee_qn, "Base") != NULL); + } + cbm_free_result(r); + PASS(); +} + /* ── Suite ─────────────────────────────────────────────────────── */ SUITE(py_lsp) { @@ -1159,4 +1211,8 @@ SUITE(py_lsp) { /* Round 4 — instance attribute typing */ RUN_TEST(pylsp_round4_instance_attribute_init); RUN_TEST(pylsp_round4_instance_attribute_class_annotation); + /* Round 5 — subscript, super().__init__, operator dunders */ + RUN_TEST(pylsp_round5_dict_subscript_value_type); + RUN_TEST(pylsp_round5_list_subscript_value_type); + RUN_TEST(pylsp_round5_super_init); }