Merge pull request #2690 from DeusData/fix/py-decorated-nested-class

fix(python): extract a decorated class nested in a class body

internal/cbm/extract_defs.c

@@ -11493,9 +11493,12 @@ static void push_nested_class_nodes(TSNode body, const CBMLangSpec *spec, wd_sta
                 wd_push(s, child, enclosing_qn);
             } else {
                 const char *ck = ts_node_type(child);
-                /* A nested class inside #if/#else is still nested here. */
+                /* A nested class inside #if/#else is still nested here, and so is
+                 * a decorated one (Python `@dataclass class Meta:` in a class
+                 * body): the decorated_definition wraps the class_definition. */
                 if (strcmp(ck, "field_declaration") == 0 ||
                     strcmp(ck, "template_declaration") == 0 || strcmp(ck, "declaration") == 0 ||
+                    (ctx->language == CBM_LANG_PYTHON && strcmp(ck, "decorated_definition") == 0) ||
                     cbm_kind_in_set(child, cbm_conditional_block_types(spec->language))) {
                     ts_nstack_push(&nc_stack, child);
                 }

tests/test_pipeline.c

@@ -15511,6 +15511,94 @@ TEST(pipeline_python_nested_helper_keeps_tests_parallel_issue1548) {
     return pyn_run_nested_helper_tests(52);
 }
 
+/* A decorated class nested in a class body (`@dataclass class Meta:` inside a
+ * class) is a nested class like an undecorated one: it gets its Class node and
+ * its methods their Method nodes. push_nested_class_nodes looked through
+ * field/template/declaration and conditional wrappers but not through Python's
+ * decorated_definition, so the decorated class and everything in it was never
+ * extracted. Both resolvers. */
+static const char pyn_decorated_nested_py[] = "def deco(cls):\n"
+                                              "    return cls\n"
+                                              "\n"
+                                              "class Outer:\n"
+                                              "    @deco\n"
+                                              "    class Decorated:\n"
+                                              "        def run(self):\n"
+                                              "            return 1\n"
+                                              "\n"
+                                              "    class Plain:\n"
+                                              "        def go(self):\n"
+                                              "            return 2\n";
+
+static bool pyn_node_exists(cbm_store_t *s, const char *project, const char *rel_qn,
+                            const char *label) {
+    char qn[512];
+    pyn_qn(qn, sizeof(qn), project, rel_qn);
+    cbm_node_t n = {0};
+    bool found = cbm_store_find_node_by_qn(s, project, qn, &n) == CBM_STORE_OK && n.label &&
+                 strcmp(n.label, label) == 0;
+    cbm_node_free_fields(&n);
+    return found;
+}
+
+static int pyn_run_decorated_nested(int fillers) {
+    char tmp[256];
+    snprintf(tmp, sizeof(tmp), "/tmp/cbm_py_decnest_XXXXXX");
+    if (!cbm_mkdtemp(tmp)) {
+        FAIL("tmpdir");
+    }
+    write_temp_file(tmp, "sd.py", pyn_decorated_nested_py);
+    for (int i = 0; i < fillers; i++) {
+        char name[64];
+        char body[128];
+        snprintf(name, sizeof(name), "filler%d.py", i);
+        snprintf(body, sizeof(body), "def filler%d():\n    return %d\n", i, i);
+        write_temp_file(tmp, name, body);
+    }
+    char *old_workers = getenv("CBM_WORKERS");
+    char *saved = old_workers ? strdup(old_workers) : NULL;
+    if (fillers > 0) {
+        cbm_setenv("CBM_WORKERS", "4", 1);
+    }
+    char db_path[512];
+    snprintf(db_path, sizeof(db_path), "%s/decnest.db", tmp);
+    cbm_pipeline_t *p = cbm_pipeline_new(tmp, db_path, CBM_MODE_FULL);
+    int rc = p ? cbm_pipeline_run(p) : -1;
+    cbm_store_t *s = rc == 0 ? cbm_store_open_path(db_path) : NULL;
+    const char *project = p ? cbm_pipeline_project_name(p) : NULL;
+    bool dec_class = s && pyn_node_exists(s, project, "sd.Outer.Decorated", "Class");
+    bool dec_method = s && pyn_node_exists(s, project, "sd.Outer.Decorated.run", "Method");
+    bool plain_class = s && pyn_node_exists(s, project, "sd.Outer.Plain", "Class");
+    bool plain_method = s && pyn_node_exists(s, project, "sd.Outer.Plain.go", "Method");
+    if (s) {
+        cbm_store_close(s);
+    }
+    if (p) {
+        cbm_pipeline_free(p);
+    }
+    if (saved) {
+        cbm_setenv("CBM_WORKERS", saved, 1);
+        free(saved);
+    } else {
+        cbm_unsetenv("CBM_WORKERS");
+    }
+    th_rmtree(tmp);
+    ASSERT_EQ(rc, 0);
+    ASSERT_TRUE(dec_class);    /* the decorated nested class is extracted */
+    ASSERT_TRUE(dec_method);   /* ... and its method */
+    ASSERT_TRUE(plain_class);  /* control: undecorated nested class */
+    ASSERT_TRUE(plain_method); /* control */
+    PASS();
+}
+
+TEST(pipeline_python_decorated_nested_class_extracted) {
+    return pyn_run_decorated_nested(0);
+}
+
+TEST(pipeline_python_decorated_nested_class_extracted_parallel) {
+    return pyn_run_decorated_nested(52);
+}
+
 TEST(pipeline_python_project) {
     /* Port of TestPipelinePythonProject */
     const char *files[] = {"main.py", "utils.py"};
@@ -25028,6 +25116,8 @@ SUITE(pipeline) {
     RUN_TEST(pipeline_python_init_function_shares_submodule_qn_parallel_issue1548);
     RUN_TEST(pipeline_python_nested_helper_keeps_tests_issue1548);
     RUN_TEST(pipeline_python_nested_helper_keeps_tests_parallel_issue1548);
+    RUN_TEST(pipeline_python_decorated_nested_class_extracted);
+    RUN_TEST(pipeline_python_decorated_nested_class_extracted_parallel);
     RUN_TEST(pipeline_header_include_target_is_independent_of_registration_order);
     RUN_TEST(pipeline_imports_multi_symbol_edges);
     RUN_TEST(pipeline_go_cross_package_call);