fix(python): resolve annotations through TYPE_CHECKING imports (#1277)

The import an annotated field's type comes from usually sits under a
module-level `if TYPE_CHECKING:` (to avoid an import cycle).
parse_python_imports walked only the module's direct statements, so that
import was never recorded and the reporter's chain
trainer.strategies.process_batch() lost its type exactly where real code
writes it.

- Extraction records the imports of a module-level `if TYPE_CHECKING:` /
  `if <module>.TYPE_CHECKING:` consequence block, flagged
  CBMImport.type_checking_only. elif/else branches stay unrecorded, as
  every other compound statement's imports do.
- One rule for every CALLS import-map builder:
  cbm_import_binds_runtime_name() (false for a flagged import). The
  sequential build_import_map applies it in its fast and slow paths; the
  shared cbm_pxc_build_import_map (= _ex(..., false)) applies it in both
  loops, so no CALLS import map ever binds an annotation-only name.
- The annotation consumers take cbm_pxc_build_import_map_ex(..., true):
  the cross-file base/field-type resolution and the Python LSP dispatch.
  The sequential dispatch always takes it; the parallel one builds that
  map only for a file that records a flagged import and otherwise passes
  its CALLS map, which is identical there by construction.
- The Python LSP binds those names into an annotation-only frame consulted
  by annotation resolution after the ordinary scope chain: a runtime call
  through such a name never resolves.

Tests: python_type_checking_imports_flagged_issue1277 (extraction),
python_type_checking_typed_field_issue1277(_parallel) (pipeline, both
venues; controls: no import_map and no LSP CALLS edge through the name).

Signed-off-by: Martin Vogel <martin.vogel.tech@gmail.com>

internal/cbm/cbm.c

@@ -402,6 +402,24 @@ void cbm_imports_push(CBMImportArray *arr, CBMArena *a, CBMImport imp) {
     arr->items[arr->count++] = imp;
 }
 
+bool cbm_imports_local_binds_runtime_name(const CBMImportArray *imports, const char *local_name) {
+    if (!imports || !local_name) {
+        return true;
+    }
+    bool seen = false;
+    for (int i = 0; i < imports->count; i++) {
+        const CBMImport *imp = &imports->items[i];
+        if (!imp->local_name || strcmp(imp->local_name, local_name) != 0) {
+            continue;
+        }
+        if (cbm_import_binds_runtime_name(imp)) {
+            return true;
+        }
+        seen = true;
+    }
+    return !seen;
+}
+
 void cbm_usages_push(CBMUsageArray *arr, CBMArena *a, CBMUsage usage) {
     GROW_ARRAY(arr, a);
     arr->items[arr->count++] = usage;

internal/cbm/cbm.h

@@ -403,6 +403,15 @@ typedef struct {
     CBMImportKind kind;      // CBM_IMPORT_KIND_DEFAULT unless the syntax names a member kind
 } CBMImport;
 
+/* True when an import binds its local name at runtime. A Python import under
+ * `if TYPE_CHECKING:` exists for annotations only: a call through that name
+ * cannot happen, so no CALLS import map may bind it (an import_map edge
+ * through it would be fabricated by construction). THE one rule every CALLS
+ * import-map builder applies. */
+static inline bool cbm_import_binds_runtime_name(const CBMImport *imp) {
+    return imp && !imp->type_checking_only;
+}
+
 /* Content-only ordered Python namespace facts; no resolved target identities. */
 #define CBM_PY_NAMESPACE_FACTS_VERSION 1u
 
@@ -1308,6 +1317,10 @@ void cbm_py_namespace_test_copy_fail_after(int successful_allocations);
 void cbm_defs_push(CBMDefArray *arr, CBMArena *a, CBMDefinition def);
 void cbm_calls_push(CBMCallArray *arr, CBMArena *a, CBMCall call);
 void cbm_imports_push(CBMImportArray *arr, CBMArena *a, CBMImport imp);
+/* True unless every import of `imports` that binds `local_name` is
+ * annotation-only (cbm_import_binds_runtime_name), so also true when no import
+ * names it: the rule for the builders that key IMPORTS edges by local name. */
+bool cbm_imports_local_binds_runtime_name(const CBMImportArray *imports, const char *local_name);
 void cbm_usages_push(CBMUsageArray *arr, CBMArena *a, CBMUsage usage);
 void cbm_throws_push(CBMThrowArray *arr, CBMArena *a, CBMThrow thr);
 void cbm_rw_push(CBMRWArray *arr, CBMArena *a, CBMReadWrite rw);

internal/cbm/extract_imports.c

@@ -321,6 +321,63 @@ static void process_py_import_from(CBMExtractCtx *ctx, TSNode node) {
     }
 }
 
+/* `TYPE_CHECKING` or `<module>.TYPE_CHECKING` (typing.TYPE_CHECKING, t.TYPE_CHECKING). */
+static bool py_is_type_checking_condition(CBMExtractCtx *ctx, TSNode cond) {
+    if (ts_node_is_null(cond)) {
+        return false;
+    }
+    TSNode name = cond;
+    if (strcmp(ts_node_type(cond), "attribute") == 0) {
+        name = ts_node_child_by_field_name(cond, TS_FIELD("attribute"));
+    }
+    if (ts_node_is_null(name) || strcmp(ts_node_type(name), "identifier") != 0) {
+        return false;
+    }
+    const char *text = cbm_node_text(ctx->arena, name, ctx->source);
+    return text && strcmp(text, "TYPE_CHECKING") == 0;
+}
+
+/* One import statement; `type_only` flags every import it records. */
+static void process_py_import_node(CBMExtractCtx *ctx, TSNode node, bool type_only) {
+    const char *kind = ts_node_type(node);
+    int before = ctx->result->imports.count;
+    if (strcmp(kind, "import_statement") == 0) {
+        process_py_import_stmt(ctx, node);
+    } else if (strcmp(kind, "import_from_statement") == 0 ||
+               strcmp(kind, "future_import_statement") == 0) {
+        // `from __future__ import annotations` is a distinct node type in
+        // tree-sitter-python but has the same shape (module + name list).
+        process_py_import_from(ctx, node);
+    }
+    for (int i = before; type_only && i < ctx->result->imports.count; i++) {
+        ctx->result->imports.items[i].type_checking_only = true;
+    }
+}
+
+/* A module-level `if TYPE_CHECKING:` holds imports a type checker sees and
+ * the interpreter never runs (#1277: an annotated field's type is usually
+ * imported exactly there, to avoid an import cycle). Record the imports of
+ * its consequence block, flagged type_checking_only; its elif/else branches
+ * are runtime alternatives and stay unrecorded, as every other compound
+ * statement's imports do. */
+static void process_py_type_checking_block(CBMExtractCtx *ctx, TSNode if_node) {
+    if (!py_is_type_checking_condition(
+            ctx, ts_node_child_by_field_name(if_node, TS_FIELD("condition")))) {
+        return;
+    }
+    TSNode body = ts_node_child_by_field_name(if_node, TS_FIELD("consequence"));
+    if (ts_node_is_null(body)) {
+        return;
+    }
+    TSTreeCursor cursor = ts_tree_cursor_new(body);
+    if (ts_tree_cursor_goto_first_child(&cursor)) {
+        do {
+            process_py_import_node(ctx, ts_tree_cursor_current_node(&cursor), true);
+        } while (ts_tree_cursor_goto_next_sibling(&cursor));
+    }
+    ts_tree_cursor_delete(&cursor);
+}
+
 static void parse_python_imports(CBMExtractCtx *ctx) {
     TSTreeCursor cursor = ts_tree_cursor_new(ctx->root);
     if (!ts_tree_cursor_goto_first_child(&cursor)) {
@@ -329,15 +386,10 @@ static void parse_python_imports(CBMExtractCtx *ctx) {
     }
     do {
         TSNode node = ts_tree_cursor_current_node(&cursor);
-        const char *kind = ts_node_type(node);
-
-        if (strcmp(kind, "import_statement") == 0) {
-            process_py_import_stmt(ctx, node);
-        } else if (strcmp(kind, "import_from_statement") == 0 ||
-                   strcmp(kind, "future_import_statement") == 0) {
-            // `from __future__ import annotations` is a distinct node type in
-            // tree-sitter-python but has the same shape (module + name list).
-            process_py_import_from(ctx, node);
+        if (strcmp(ts_node_type(node), "if_statement") == 0) {
+            process_py_type_checking_block(ctx, node);
+        } else {
+            process_py_import_node(ctx, node, false);
         }
     } while (ts_tree_cursor_goto_next_sibling(&cursor));
     ts_tree_cursor_delete(&cursor);

internal/cbm/lsp/py_lsp.c

@@ -4278,6 +4278,11 @@ static const CBMType *py_resolve_annotation(PyLSPContext *ctx, const char *ann)
     const CBMType *t = cbm_scope_lookup(ctx->current_scope, ann);
     if (!cbm_type_is_unknown(t))
         return t;
+    if (ctx->annotation_scope) {
+        t = cbm_scope_lookup(ctx->annotation_scope, ann);
+        if (!cbm_type_is_unknown(t))
+            return t;
+    }
     if (ctx->module_qn) {
         const char *qn = cbm_arena_sprintf(ctx->arena, "%s.%s", ctx->module_qn, ann);
         const CBMRegisteredType *rt = cbm_registry_lookup_type(ctx->registry, qn);
@@ -5360,6 +5365,27 @@ static void py_replay_import_statement(PyLSPContext *ctx, TSNode stmt, unsigned
     }
 }
 
+/* The consequence block of a module-level `if TYPE_CHECKING:` /
+ * `if <module>.TYPE_CHECKING:` statement, else a null node. Mirrors the
+ * extraction's py_is_type_checking_condition (extract_imports.c). */
+static TSNode py_type_checking_body(PyLSPContext *ctx, TSNode stmt) {
+    TSNode none = {0};
+    if (strcmp(ts_node_type(stmt), "if_statement") != 0)
+        return none;
+    TSNode cond = ts_node_child_by_field_name(stmt, "condition", 9);
+    if (ts_node_is_null(cond))
+        return none;
+    TSNode name = cond;
+    if (strcmp(ts_node_type(cond), "attribute") == 0)
+        name = ts_node_child_by_field_name(cond, "attribute", 9);
+    if (ts_node_is_null(name) || strcmp(ts_node_type(name), "identifier") != 0)
+        return none;
+    char *text = py_node_text(ctx, name);
+    if (!text || strcmp(text, "TYPE_CHECKING") != 0)
+        return none;
+    return ts_node_child_by_field_name(stmt, "consequence", 11);
+}
+
 static bool py_expression_is_annotation_only_assignment(TSNode statement) {
     if (ts_node_named_child_count(statement) != 1)
         return false;
@@ -5444,6 +5470,29 @@ void py_lsp_process_file(PyLSPContext *ctx, TSNode root) {
              * Invalidate only syntactic binding targets at the control-flow
              * join; a later unconditional statement may restore exact proof. */
             py_invalidate_possible_bindings(ctx, c, 0);
+            /* `if TYPE_CHECKING:` holds the imports annotations are written
+             * against (#1277). Bind them into the annotation-only frame: an
+             * annotation resolves through them, a runtime call never does (the
+             * interpreter never runs the block). Extraction recorded them
+             * flagged type_checking_only, so they reach this walk through the
+             * LSP import map only, never a CALLS import map. */
+            TSNode tc_body = py_type_checking_body(ctx, c);
+            if (!ts_node_is_null(tc_body)) {
+                if (!ctx->annotation_scope)
+                    ctx->annotation_scope = cbm_scope_push(ctx->arena, NULL);
+                CBMScope *runtime_scope = ctx->current_scope;
+                PyKids tk = py_kids(ctx->arena, tc_body);
+                for (uint32_t t = 0; ctx->annotation_scope && t < tk.n && !ctx->eval_failure; t++) {
+                    TSNode stmt = py_kid(&tk, t);
+                    const char *sk = ts_node_type(stmt);
+                    if (strcmp(sk, "import_statement") == 0 ||
+                        strcmp(sk, "import_from_statement") == 0) {
+                        ctx->current_scope = ctx->annotation_scope;
+                        py_replay_import_statement(ctx, stmt, consumed_imports);
+                        ctx->current_scope = runtime_scope;
+                    }
+                }
+            }
         }
     }
     // Pass 2: top-level calls (rare) and nested definitions.

internal/cbm/lsp/py_lsp.h

@@ -49,6 +49,10 @@ typedef struct {
      * cbm_registry_*_for_update, never into anything behind ->fallback. */
     CBMTypeRegistry *registry_head;
     CBMScope *current_scope;
+    // Names a module-level `if TYPE_CHECKING:` imports (#1277): consulted by
+    // annotation resolution only, after the ordinary scope chain misses, so a
+    // runtime call through such a name never resolves. NULL when none.
+    CBMScope *annotation_scope;
 
     // Import map: local_name -> module_qn (arena-allocated, NULL-terminated).
     // Built from CBMFileResult.imports.

src/pipeline/pass_calls.c

@@ -141,7 +141,8 @@ static int build_import_map(cbm_pipeline_ctx_t *ctx, const char *rel_path,
 
         for (int i = 0; i < result->imports.count; i++) {
             const CBMImport *imp = &result->imports.items[i];
-            if (!imp->local_name || !imp->local_name[0] || !imp->module_path) {
+            if (!imp->local_name || !imp->local_name[0] || !imp->module_path ||
+                !cbm_import_binds_runtime_name(imp)) {
                 continue;
             }
             const cbm_gbuf_node_t *target =
@@ -193,6 +194,10 @@ static int build_import_map(cbm_pipeline_ctx_t *ctx, const char *rel_path,
             continue;
         }
         char *key = extract_local_name_from_json(e->properties_json);
+        if (key && !cbm_imports_local_binds_runtime_name(result ? &result->imports : NULL, key)) {
+            free(key); /* annotation-only import (#1277): no CALLS binding */
+            key = NULL;
+        }
         if (key) {
             keys[count] = key;
             vals[count] = target->qualified_name;

src/pipeline/pass_lsp_cross.c

@@ -876,9 +876,11 @@ CBMLSPDef *cbm_pxc_collect_all_defs(const cbm_pipeline_ctx_t *ctx, CBMArena *are
         if (ctx && ctx->registry && pxc_lang_resolves_base_qns(files[fi].language)) {
             bs.reg = ctx->registry;
             bs.gbuf = ctx->gbuf;
-            cbm_pxc_build_import_map(ctx->gbuf, project_name, files[fi].rel_path,
-                                     files[fi].language, fr, &bs.imp_keys, &bs.imp_vals,
-                                     &bs.imp_count);
+            /* Bases and field annotations resolve through annotation-only
+             * imports too: `if TYPE_CHECKING:` exists for them (#1277). */
+            cbm_pxc_build_import_map_ex(ctx->gbuf, project_name, files[fi].rel_path,
+                                        files[fi].language, fr, true, &bs.imp_keys, &bs.imp_vals,
+                                        &bs.imp_count);
         }
         for (int di = 0; di < fr->defs.count; di++) {
             if (pxc_build_lsp_def(arena, &fr->defs.items[di], def_modules[fi], namespace_name,
@@ -1068,10 +1070,26 @@ static char *pxc_python_import_from_metadata(const cbm_gbuf_t *gbuf, const char
 
 /* Build per-file import map (local_name -> semantic import QN) from gbuf
  * IMPORTS edges. Both pipeline drivers call this implementation. Returns 0
- * with *out_count = 0 when the file has no IMPORTS edges. */
+ * with *out_count = 0 when the file has no IMPORTS edges. The CALLS map. */
 int cbm_pxc_build_import_map(const cbm_gbuf_t *gbuf, const char *project_name, const char *rel_path,
                              CBMLanguage lang, const CBMFileResult *result, const char ***out_keys,
                              const char ***out_vals, int *out_count) {
+    return cbm_pxc_build_import_map_ex(gbuf, project_name, rel_path, lang, result, false, out_keys,
+                                       out_vals, out_count);
+}
+
+bool cbm_pxc_result_has_type_only_import(const CBMFileResult *result) {
+    for (int i = 0; result && i < result->imports.count; i++) {
+        if (!cbm_import_binds_runtime_name(&result->imports.items[i]))
+            return true;
+    }
+    return false;
+}
+
+int cbm_pxc_build_import_map_ex(const cbm_gbuf_t *gbuf, const char *project_name,
+                                const char *rel_path, CBMLanguage lang, const CBMFileResult *result,
+                                bool include_type_only, const char ***out_keys,
+                                const char ***out_vals, int *out_count) {
     *out_keys = NULL;
     *out_vals = NULL;
     *out_count = 0;
@@ -1129,6 +1147,11 @@ int cbm_pxc_build_import_map(const cbm_gbuf_t *gbuf, const char *project_name, c
             continue;
         memcpy(local, start, n);
         local[n] = '\0';
+        if (!include_type_only &&
+            !cbm_imports_local_binds_runtime_name(result ? &result->imports : NULL, local)) {
+            free(local); /* annotation-only import (#1277): no CALLS binding */
+            continue;
+        }
         char *value = pxc_import_value_qn(lang, result, local, target);
         if (!value) {
             free(local);
@@ -1146,6 +1169,7 @@ int cbm_pxc_build_import_map(const cbm_gbuf_t *gbuf, const char *project_name, c
     for (int i = 0; i < metadata_count; i++) {
         const CBMImport *imp = &result->imports.items[i];
         if (!imp->local_name || !imp->local_name[0] || !imp->module_path ||
+            (!include_type_only && !cbm_import_binds_runtime_name(imp)) ||
             pxc_import_map_has_local(keys, count, imp->local_name)) {
             continue;
         }
@@ -2212,8 +2236,10 @@ int cbm_pipeline_pass_lsp_cross(cbm_pipeline_ctx_t *ctx, const cbm_file_info_t *
         const char **imp_keys = NULL;
         const char **imp_vals = NULL;
         int imp_count = 0;
-        cbm_pxc_build_import_map(ctx->gbuf, ctx->project_name, files[i].rel_path, lang, cache[i],
-                                 &imp_keys, &imp_vals, &imp_count);
+        /* The LSP map: annotation-only imports bind for annotations (#1277).
+         * MUST match the LSP map of pass_parallel.c. */
+        cbm_pxc_build_import_map_ex(ctx->gbuf, ctx->project_name, files[i].rel_path, lang, cache[i],
+                                    true, &imp_keys, &imp_vals, &imp_count);
 
         /* Journal around the resolve: a hang here must be attributed to THIS
          * file, not to a stale extraction marker (the innocent-quarantine

src/pipeline/pass_lsp_cross.h

@@ -83,6 +83,17 @@ void cbm_pxc_ts_modes(CBMLanguage lang, const char *rel_path, bool *out_js, bool
 int cbm_pxc_build_import_map(const cbm_gbuf_t *gbuf, const char *project_name, const char *rel_path,
                              CBMLanguage lang, const CBMFileResult *result, const char ***out_keys,
                              const char ***out_vals, int *out_count);
+/* As cbm_pxc_build_import_map (which is include_type_only = false, the CALLS
+ * map), but with include_type_only also binding annotation-only imports
+ * (Python `if TYPE_CHECKING:`), for the consumers that resolve annotations:
+ * the Python LSP and the cross-file base/field-type resolution (#1277). */
+int cbm_pxc_build_import_map_ex(const cbm_gbuf_t *gbuf, const char *project_name,
+                                const char *rel_path, CBMLanguage lang, const CBMFileResult *result,
+                                bool include_type_only, const char ***out_keys,
+                                const char ***out_vals, int *out_count);
+/* True when the file records an annotation-only import (the LSP map then
+ * differs from the CALLS map and must be built separately). */
+bool cbm_pxc_result_has_type_only_import(const CBMFileResult *result);
 
 void cbm_pxc_free_import_map(const char **keys, const char **vals, int count);
 

src/pipeline/pass_parallel.c

@@ -4323,10 +4323,27 @@ static void resolve_worker(int worker_id, void *ctx_ptr) {
                  * this file, not to a stale extraction marker. */
                 cbm_index_mark_start(rel);
                 bool had_error = result->has_error;
+                /* The LSP map also binds annotation-only imports (#1277);
+                 * only a file that records one needs its own map, every
+                 * other file passes the CALLS map unchanged. MUST match the
+                 * sequential dispatch in pass_lsp_cross.c. */
+                const char **lsp_keys = NULL;
+                const char **lsp_vals = NULL;
+                int lsp_count = 0;
+                bool own_lsp_map = cbm_pxc_result_has_type_only_import(result);
+                if (own_lsp_map) {
+                    cbm_pxc_build_import_map_ex(rc->main_gbuf, rc->project_name, rel, lang, result,
+                                                true, &lsp_keys, &lsp_vals, &lsp_count);
+                }
                 cbm_pxc_dispatch_file_with_php_bindings(
                     lang, result, lsp_source, lsp_source_len, rel, def_module, rc->cross_registries,
-                    rc->module_def_index, rc->all_defs, rc->def_count, imp_keys, imp_vals,
-                    imp_count, pp_rust_shared_registry_get, rc, &php_vendor);
+                    rc->module_def_index, rc->all_defs, rc->def_count,
+                    own_lsp_map ? lsp_keys : imp_keys, own_lsp_map ? lsp_vals : imp_vals,
+                    own_lsp_map ? lsp_count : imp_count, pp_rust_shared_registry_get, rc,
+                    &php_vendor);
+                if (own_lsp_map) {
+                    cbm_pxc_free_import_map(lsp_keys, lsp_vals, lsp_count);
+                }
                 cbm_index_mark_done(rel);
                 if (!had_error && result->has_error && result->error_msg &&
                     (strcmp(result->error_msg, CBM_C_LSP_MEMO_ERROR) == 0 ||