trending/codebase-memory-mcp
1
0
Fork 0
mirror of https://github.com/DeusData/codebase-memory-mcp.git synced 2026-10-26 06:43:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(py_lsp): TypedDict subscript + early-return narrowing + match-sequence

Browse Source 
Round 9 closes 2 of the 3 remaining stress gaps from Round 8:

TypedDict literal-key subscript:
- d["foo"] where d: TD with class-body annotation `foo: Foo` returns
  the field's type. Routes through the existing class-body-annotation
  field registry (Round 4): TypedDict subclasses already register their
  fields the same way. The subscript branch in py_eval_expr_type now
  detects NAMED container + literal-string key and looks up the field
  via py_lookup_field. Strips both single and double quotes.

Early-return narrowing (post-block flow):
- if not isinstance(x, T): return / raise / break / continue —
  applies the *positive* narrow to the enclosing scope after the
  if-block when the consequence terminates.
- if x is None: return — narrows x to non-None afterwards.
- py_block_terminates detects return/raise/break/continue as the last
  statement of the body block. Subsequent code in the function sees
  the narrowed type.

Match sequence pattern (tentative):
- list_pattern / tuple_pattern handler added: bind capture-pattern
  elements to the subject's iterable element type; bind splat-pattern
  rest to list[elem]. Also walks more wrapper levels (case_pattern ->
  pattern -> _simple_pattern). Test still flags as KNOWN GAP because
  tree-sitter's pattern-wrapper structure for the bracket-form `[a, b]`
  in match/case differs in some grammar versions; documented as
  remaining gap pending grammar inspection.

Stress test suite: 22/22 PASS. Hard-asserted patterns: namedtuple,
typeddict, protocol, abc, with-as, except-as, isinstance-early-return,
tuple-unpack, dict.items, slice, decorator-factory, property-setter,
Self chain, diamond MRO, recursive types, nested closure, long method
chains, generator delegation, async-gen for. Documented gaps:
function-as-dict-value, match sequence pattern, lambda inference.

Bench unchanged at 100% / 11.02 ms; total 2920 / 0.
This commit is contained in:
Martin Vogel
2026-05-09 12:48:11 +02:00
parent f35745c79f
commit 4ccd08c6de
2 changed files with 161 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/cbm/lsp/py_lsp.c
+153 -7
View File
@@ -840,11 +840,40 @@ 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).
// ...])[int] -> A (union when index isn't a literal int). NAMED that
// resolves to a TypedDict-tagged registered type with a literal-string
// key returns the field's type.
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);
// TypedDict literal-key subscript: d["foo"] where d: TD and TD has
// class-body annotation `foo: Foo`. Field is registered via
// py_register_instance_field.
if (container && container->kind == CBM_TYPE_NAMED) {
TSNode sub = ts_node_child_by_field_name(node, "subscript", 9);
if (!ts_node_is_null(sub)) {
const char* sk = ts_node_type(sub);
if (strcmp(sk, "string") == 0) {
char* lit = py_node_text(ctx, sub);
if (lit) {
// Strip surrounding quotes.
size_t llen = strlen(lit);
if (llen >= 2 && (lit[0] == '"' || lit[0] == '\'') &&
lit[llen - 1] == lit[0]) {
lit[llen - 1] = '\0';
lit++;
}
const CBMType* field = py_lookup_field(ctx,
container->data.named.qualified_name, lit);
if (field) return field;
}
}
}
return cbm_type_unknown();
}
if (!container || container->kind != CBM_TYPE_TEMPLATE) {
return cbm_type_unknown();
}
@@ -1437,8 +1466,25 @@ static void py_bind_walrus_in(PyLSPContext* ctx, TSNode node) {
}
}
/* Detect whether a block ends with `return`, `raise`, `break`, or
* `continue` — meaning execution doesn't fall through past the if-block.
* Used for early-return narrowing of the negated condition. */
static bool py_block_terminates(TSNode block) {
if (ts_node_is_null(block)) return false;
uint32_t bnc = ts_node_named_child_count(block);
if (bnc == 0) return false;
TSNode last = ts_node_named_child(block, bnc - 1);
const char* k = ts_node_type(last);
return strcmp(k, "return_statement") == 0 ||
strcmp(k, "raise_statement") == 0 ||
strcmp(k, "break_statement") == 0 ||
strcmp(k, "continue_statement") == 0;
}
/* Walk if_statement: push scope for consequence body when narrowing
* applies, then recurse. */
* applies, then recurse. If the consequence terminates (return/raise),
* apply the negated narrow to the *enclosing* scope so subsequent
* statements see the narrowed type. */
static void py_walk_if_statement(PyLSPContext* ctx, TSNode if_node) {
TSNode cond = ts_node_child_by_field_name(if_node, "condition", 9);
TSNode body = ts_node_child_by_field_name(if_node, "consequence", 11);
@@ -1450,6 +1496,39 @@ static void py_walk_if_statement(PyLSPContext* ctx, TSNode if_node) {
// Always evaluate the condition (may bind via walrus, currently a no-op).
py_resolve_calls_in(ctx, cond);
// Recognize negated isinstance / is-None patterns for early-return
// narrowing.
bool neg_isinstance = false;
char* neg_name = NULL;
const CBMType* neg_type = NULL;
int neg_none_kind = 0;
char* neg_none_var = NULL;
if (!ts_node_is_null(cond) &&
strcmp(ts_node_type(cond), "not_operator") == 0 &&
ts_node_named_child_count(cond) > 0) {
TSNode inner = ts_node_named_child(cond, 0);
if (strcmp(ts_node_type(inner), "call") == 0) {
char* nm = NULL;
const CBMType* ty = NULL;
if (py_match_isinstance(ctx, inner, &nm, &ty)) {
neg_isinstance = true;
neg_name = nm;
neg_type = ty;
}
}
// `if not (x is None):` -> equivalent to `if x is not None`
// (rare, the next branch handles `is not None` directly).
}
// `if x is None: return` — after the if block, x is non-None.
{
char* nv = NULL;
int nk = py_match_is_none(ctx, cond, &nv);
if (nk == -1) { // x is None
neg_none_kind = -1;
neg_none_var = nv;
}
}
// Consequence: maybe narrow.
if (!ts_node_is_null(body)) {
CBMScope* saved = ctx->current_scope;
@@ -1483,6 +1562,24 @@ static void py_walk_if_statement(PyLSPContext* ctx, TSNode if_node) {
if (!ts_node_is_null(alt)) {
py_resolve_calls_in(ctx, alt);
}
// Early-return narrowing: if the consequence block terminates and the
// condition negates a type guard, apply the *positive* narrow to the
// enclosing scope so subsequent statements see the narrowed type.
if (!ts_node_is_null(body) && py_block_terminates(body)) {
if (neg_isinstance && neg_name && neg_type) {
cbm_scope_bind(ctx->current_scope, neg_name, neg_type);
}
if (neg_none_kind == -1 && neg_none_var) {
// `if x is None: return` -> narrow x to non-None afterwards.
const CBMType* current = cbm_scope_lookup(ctx->current_scope,
neg_none_var);
if (current && !cbm_type_is_unknown(current)) {
const CBMType* narrowed = py_strip_none(ctx, current);
cbm_scope_bind(ctx->current_scope, neg_none_var, narrowed);
}
}
}
}
static void py_resolve_calls_in(PyLSPContext* ctx, TSNode node) {
@@ -1536,12 +1633,20 @@ static void py_resolve_calls_in(PyLSPContext* ctx, TSNode node) {
// The pattern field is a `case_pattern` wrapper around the
// actual pattern type (class_pattern, sequence_pattern, etc).
// Unwrap one level if needed.
// Some grammar versions wrap further in `pattern`. Walk
// through both wrappers.
TSNode actual_pattern = pattern;
if (!ts_node_is_null(pattern) &&
strcmp(ts_node_type(pattern), "case_pattern") == 0 &&
ts_node_named_child_count(pattern) > 0) {
actual_pattern = ts_node_named_child(pattern, 0);
for (int unwrap = 0; unwrap < 3; unwrap++) {
if (ts_node_is_null(actual_pattern)) break;
const char* apk = ts_node_type(actual_pattern);
if ((strcmp(apk, "case_pattern") == 0 ||
strcmp(apk, "pattern") == 0 ||
strcmp(apk, "_simple_pattern") == 0) &&
ts_node_named_child_count(actual_pattern) > 0) {
actual_pattern = ts_node_named_child(actual_pattern, 0);
continue;
}
break;
}
// Class pattern: case Foo(a, b): -> bind subject to NAMED(Foo),
// bind a / b to UNKNOWN (field-type extraction is a v1.1 task).
@@ -1568,6 +1673,47 @@ static void py_resolve_calls_in(PyLSPContext* ctx, TSNode node) {
}
}
}
// Sequence pattern: `case [head, *tail]:` / `case (a, b):`
// — for each capture pattern element, bind to the
// subject's iterable element type. Star-rest binds to
// list[elem].
if (subject_name && !ts_node_is_null(actual_pattern) &&
(strcmp(ts_node_type(actual_pattern), "list_pattern") == 0 ||
strcmp(ts_node_type(actual_pattern), "tuple_pattern") == 0)) {
const CBMType* subj_t = cbm_scope_lookup(ctx->current_scope,
subject_name);
const CBMType* elem_t = py_iterable_element_type(ctx, subj_t);
uint32_t pn = ts_node_named_child_count(actual_pattern);
for (uint32_t j = 0; j < pn; j++) {
TSNode sub = ts_node_named_child(actual_pattern, j);
const char* sk = ts_node_type(sub);
if (strcmp(sk, "capture_pattern") == 0) {
if (ts_node_named_child_count(sub) > 0) {
TSNode id = ts_node_named_child(sub, 0);
char* nm = py_node_text(ctx, id);
if (nm && elem_t) {
cbm_scope_bind(ctx->current_scope, nm, elem_t);
}
}
} else if (strcmp(sk, "splat_pattern") == 0 ||
strcmp(sk, "list_splat_pattern") == 0 ||
strcmp(sk, "star_pattern") == 0) {
if (ts_node_named_child_count(sub) > 0) {
TSNode id = ts_node_named_child(sub, 0);
if (strcmp(ts_node_type(id), "capture_pattern") == 0 &&
ts_node_named_child_count(id) > 0) {
id = ts_node_named_child(id, 0);
}
char* nm = py_node_text(ctx, id);
if (nm && elem_t) {
const CBMType* lst = cbm_type_template(
ctx->arena, "list", &elem_t, 1);
cbm_scope_bind(ctx->current_scope, nm, lst);
}
}
}
}
}
if (!ts_node_is_null(conseq)) {
py_resolve_calls_in(ctx, conseq);
tests/test_py_lsp_stress.c
+8 -7
View File
@@ -80,9 +80,7 @@ TEST(stress_typeddict_subscript) {
"def use(d: TD):\n"
" return d['foo'].method()\n");
ASSERT_NOT_NULL(r);
/* TypedDict["literal"] should narrow to the field type — gap if fails */
int idx = find_resolved(r, "use", "method");
if (idx < 0) printf(" KNOWN GAP: TypedDict['literal'] narrowing\n");
ASSERT_GTE(require_resolved(r, "use", "method"), 0);
cbm_free_result(r);
PASS();
}
@@ -168,7 +166,7 @@ TEST(stress_except_as_binding) {
TEST(stress_isinstance_else_negative_narrow) {
/* if not isinstance(x, Foo): return; x.method()
* Pyright narrows x to Foo after the early return. We don't model this. */
* Narrows x to Foo after the early return via py_block_terminates. */
CBMFileResult *r = extract_py(
"class Foo:\n"
" def method(self):\n"
@@ -178,8 +176,7 @@ TEST(stress_isinstance_else_negative_narrow) {
" return None\n"
" return x.method()\n");
ASSERT_NOT_NULL(r);
int idx = find_resolved(r, "use", "method");
if (idx < 0) printf(" KNOWN GAP: post-early-return narrowing\n");
ASSERT_GTE(require_resolved(r, "use", "method"), 0);
cbm_free_result(r);
PASS();
}
@@ -390,7 +387,11 @@ TEST(stress_nested_closure) {
}
TEST(stress_match_sequence_pattern) {
/* match items: case [head, *tail]: head.method() */
/* match items: case [head, *tail]: head.method()
* Tree-sitter Python's pattern wrapper layout for sequence patterns
* doesn't expose list_pattern as our case-pattern unwrap finds it;
* the bracket form `[head, *tail]` parses through case_pattern ->
* (a wrapper that varies by grammar version). Documented as gap. */
CBMFileResult *r = extract_py(
"class Foo:\n"
" def method(self):\n"