From e8c99d10d8b9c49dd2cea09948322bec21634d8e Mon Sep 17 00:00:00 2001 From: GertL Date: Sun, 14 Jun 2026 16:12:14 +0200 Subject: [PATCH 1/5] Enhance Python extractor with Django support and update tests --- internal/extractors/pythonextractor/python.go | 319 ++++++++- .../extractors/pythonextractor/python_test.go | 633 +++++++++++++++++- 2 files changed, 902 insertions(+), 50 deletions(-) diff --git a/internal/extractors/pythonextractor/python.go b/internal/extractors/pythonextractor/python.go index 304bb47..ae17526 100644 --- a/internal/extractors/pythonextractor/python.go +++ b/internal/extractors/pythonextractor/python.go @@ -43,10 +43,10 @@ func (e *PythonExtractor) Detect(repoPath string) (bool, error) { // Subdirectory search (up to 3 levels deep) — handles monorepos. subMarkers := map[string]bool{ - "pyproject.toml": true, - "setup.py": true, + "pyproject.toml": true, + "setup.py": true, "requirements.txt": true, - "Pipfile": true, + "Pipfile": true, } found := false _ = filepath.WalkDir(repoPath, func(path string, d fs.DirEntry, err error) error { @@ -73,6 +73,7 @@ func (e *PythonExtractor) Detect(repoPath string) (bool, error) { func (e *PythonExtractor) Extract(ctx context.Context, repoPath string, files []string) ([]facts.Fact, error) { var allFacts []facts.Fact modules := make(map[string]bool) + isDjango := detectDjango(repoPath) for _, relFile := range files { select { @@ -92,7 +93,7 @@ func (e *PythonExtractor) Extract(ctx context.Context, repoPath string, files [] continue } - fileFacts := extractFile(f, relFile) + fileFacts := extractFile(f, relFile, isDjango) f.Close() allFacts = append(allFacts, fileFacts...) @@ -135,6 +136,52 @@ var ( // tableNameRe matches SQLAlchemy __tablename__ assignments. Group: (table). tableNameRe = regexp.MustCompile(`^\s*__tablename__\s*=\s*["']([^"']+)["']`) + + // decoratorRe captures the full decorator name for structural prop detection. + // Group: (name) e.g. "staticmethod", "app.task". + decoratorRe = regexp.MustCompile(`^\s*@([\w.]+)`) + + // returnTypeRe extracts return type from a single-line def signature. + // e.g. "def foo(x: int) -> Optional[str]:" + returnTypeRe = regexp.MustCompile(`\)\s*->\s*([\w\[\], |.]+?)\s*:`) + + // returnTypeClosingRe matches the closing paren of a multi-line def with + // a return type annotation. e.g. " ) -> Optional[str]:" + returnTypeClosingRe = regexp.MustCompile(`^\s*\)\s*->\s*([\w\[\], |.]+?)\s*:`) + + // apiViewRe matches Django REST Framework @api_view decorators. + // Group: (methods_list) — bracket contents, e.g. "'GET', 'POST'" + apiViewRe = regexp.MustCompile(`^\s*@(?:[\w.]*\.)?api_view\s*\(\s*\[([^\]]+)\]`) + + // httpMethodWordRe extracts uppercase HTTP method tokens from an api_view list. + httpMethodWordRe = regexp.MustCompile(`[A-Z]+`) + + // urlPathRe matches Django path() and re_path() calls in urls.py. + // Groups: (url_path, view_ref) + urlPathRe = regexp.MustCompile(`(?:re_)?path\s*\(\s*r?["']([^"']+)["']\s*,\s*([\w.]+)`) +) + +// Django class base sets used to classify models, views, and serializers. +var ( + djangoModelBases = map[string]bool{ + "Model": true, "AbstractModel": true, "MPTTModel": true, + "TimeStampedModel": true, "UUIDModel": true, "PolymorphicModel": true, + } + + djangoCBVBases = map[string]bool{ + "View": true, "APIView": true, "GenericAPIView": true, + "ListAPIView": true, "CreateAPIView": true, "RetrieveAPIView": true, + "UpdateAPIView": true, "DestroyAPIView": true, "ListCreateAPIView": true, + "RetrieveUpdateDestroyAPIView": true, "ViewSet": true, "ModelViewSet": true, + "ReadOnlyModelViewSet": true, "TemplateView": true, "DetailView": true, + "ListView": true, "CreateView": true, "UpdateView": true, "DeleteView": true, + "FormView": true, "RedirectView": true, + } + + djangoSerializerBases = map[string]bool{ + "Serializer": true, "ModelSerializer": true, + "HyperlinkedModelSerializer": true, "ListSerializer": true, + } ) // scopeEntry tracks a class nesting level with its indentation. @@ -153,24 +200,28 @@ type pendingRoute struct { } // extractFile parses a single Python file and returns facts. -func extractFile(f *os.File, relFile string) []facts.Fact { +func extractFile(f *os.File, relFile string, isDjango bool) []facts.Fact { var result []facts.Fact dir := filepath.Dir(relFile) // Python modules are file-based; strip .py to form the module prefix used in // symbol names (e.g. "app/models/order" for "app/models/order.py"). - // This avoids name collisions between classes in different files of the same - // directory. module := strings.TrimSuffix(relFile, ".py") + isURLsFile := isDjango && filepath.Base(relFile) == "urls.py" scanner := bufio.NewScanner(f) scanner.Buffer(make([]byte, 0, 256*1024), 1024*1024) var ( - lineNum int - scopeStack []scopeEntry - pendingRoutes []pendingRoute - inDocstring bool - docstringQuote string // `"""` or `'''` + lineNum int + scopeStack []scopeEntry + pendingRoutes []pendingRoute + pendingDecorators []string // decorator names since last def/class + pendingApiViewMethods []string // HTTP methods from @api_view + decoratorParenDepth int // open-paren depth inside a multi-line decorator arg list + pendingFuncProps map[string]any // props of last emitted func (for return-type backfill) + pendingFuncLine int + inDocstring bool + docstringQuote string // `"""` or `'''` ) for scanner.Scan() { @@ -209,15 +260,28 @@ func extractFile(f *os.File, relFile string) []facts.Fact { indent := lineIndent(line) // Pop scope entries that are at the same or deeper indentation level. - // This handles returning to outer scope when indentation decreases. scopeStack = popScopes(scopeStack, indent) + // Expire pending return-type backfill after 20 lines. + if pendingFuncProps != nil && lineNum-pendingFuncLine > 20 { + pendingFuncProps = nil + } + + // Check for return type on the closing line of a multi-line function + // signature, e.g. " ) -> Optional[str]:". The props map is shared with + // the already-emitted fact, so updating it here updates the fact in-place. + if pendingFuncProps != nil { + if rt := returnTypeClosingRe.FindStringSubmatch(line); rt != nil { + pendingFuncProps["return_type"] = strings.TrimSpace(rt[1]) + pendingFuncProps = nil + } + } + // Class declaration. if m := classRe.FindStringSubmatch(line); m != nil { // m[1]=indent, m[2]=name, m[3]=bases (may be empty) name := m[2] basesStr := strings.TrimSpace(m[3]) - qualName := buildQualName(module, scopeStack, name) props := map[string]any{ @@ -229,14 +293,50 @@ func extractFile(f *os.File, relFile string) []facts.Fact { {Kind: facts.RelDeclares, Target: dir}, } + bases := splitBases(basesStr) + // Emit RelImplements for each base class. - if basesStr != "" { - for _, base := range splitBases(basesStr) { - if base != "" { - rels = append(rels, facts.Relation{ - Kind: facts.RelImplements, - Target: base, + for _, base := range bases { + if base != "" { + rels = append(rels, facts.Relation{ + Kind: facts.RelImplements, + Target: base, + }) + } + } + + // Django-specific class classification. + if isDjango { + for _, base := range bases { + bn := lastComponent(base) + if djangoModelBases[bn] { + // Emit KindStorage for the Django-inferred table name. + result = append(result, facts.Fact{ + Kind: facts.KindStorage, + Name: camelToSnake(name), + File: relFile, + Line: lineNum, + Props: map[string]any{ + "storage_kind": "table", + "framework": "django", + "language": "python", + "class": qualName, + }, + Relations: []facts.Relation{ + {Kind: facts.RelDeclares, Target: dir}, + }, }) + break + } + if djangoCBVBases[bn] { + props["django_component"] = "view" + props["framework"] = "django" + break + } + if djangoSerializerBases[bn] { + props["django_component"] = "serializer" + props["framework"] = "django" + break } } } @@ -256,9 +356,12 @@ func extractFile(f *os.File, relFile string) []facts.Fact { indent: indent, }) - // A class declaration closes any pending route (decorators above a class - // are not route handlers). + // A class declaration closes all pending state. pendingRoutes = nil + pendingDecorators = nil + pendingApiViewMethods = nil + decoratorParenDepth = 0 + pendingFuncProps = nil continue } @@ -289,6 +392,23 @@ func extractFile(f *os.File, relFile string) []facts.Fact { props["async"] = true } + // Apply structural props from pending decorators. + for _, dec := range pendingDecorators { + applyDecoratorProps(props, dec) + } + pendingDecorators = nil + decoratorParenDepth = 0 + + // Extract return type from the current line (single-line signature). + if rt := returnTypeRe.FindStringSubmatch(line); rt != nil { + props["return_type"] = strings.TrimSpace(rt[1]) + pendingFuncProps = nil + } else { + // Multi-line signature: backfill return type when closing line appears. + pendingFuncProps = props + pendingFuncLine = lineNum + } + fact := facts.Fact{ Kind: facts.KindSymbol, Name: fullName, @@ -300,7 +420,7 @@ func extractFile(f *os.File, relFile string) []facts.Fact { }, } - // If there are pending route decorators, emit route facts now. + // Emit FastAPI route facts. for _, pr := range pendingRoutes { result = append(result, facts.Fact{ Kind: facts.KindRoute, @@ -318,6 +438,24 @@ func extractFile(f *os.File, relFile string) []facts.Fact { } pendingRoutes = nil + // Emit Django @api_view route facts. + for _, method := range pendingApiViewMethods { + result = append(result, facts.Fact{ + Kind: facts.KindRoute, + Name: method + " (view) " + fullName, + File: relFile, + Line: lineNum, + Props: map[string]any{ + "http_method": method, + "path": "", + "handler": fullName, + "framework": "django", + "language": "python", + }, + }) + } + pendingApiViewMethods = nil + result = append(result, fact) continue } @@ -334,14 +472,61 @@ func extractFile(f *os.File, relFile string) []facts.Fact { continue } - // Non-route decorator — keep any pending routes (multiple decorators on - // the same function are allowed), but don't reset them here. + // Any decorator line (@...). if strings.HasPrefix(trimmed, "@") { + // Django @api_view(['GET', 'POST']) — parse HTTP methods. + if isDjango { + if m := apiViewRe.FindStringSubmatch(line); m != nil { + methods := httpMethodWordRe.FindAllString(m[1], -1) + pendingApiViewMethods = append(pendingApiViewMethods, methods...) + decoratorParenDepth = 0 // api_view always single-line + continue + } + } + + // Capture decorator name for structural symbol props. + if m := decoratorRe.FindStringSubmatch(line); m != nil { + pendingDecorators = append(pendingDecorators, m[1]) + } + // Track open parens so multi-line decorator args don't clear pending state. + decoratorParenDepth = strings.Count(line, "(") - strings.Count(line, ")") + if decoratorParenDepth < 0 { + decoratorParenDepth = 0 + } continue } - // Any non-decorator, non-def line clears pending routes. + // Inside a multi-line decorator argument list — update depth and wait. + if decoratorParenDepth > 0 { + decoratorParenDepth += strings.Count(line, "(") - strings.Count(line, ")") + if decoratorParenDepth < 0 { + decoratorParenDepth = 0 + } + continue + } + + // Any non-decorator, non-def line clears all pending state. pendingRoutes = nil + pendingDecorators = nil + pendingApiViewMethods = nil + + // Django URL patterns in urls.py files. + if isURLsFile { + if m := urlPathRe.FindStringSubmatch(line); m != nil { + result = append(result, facts.Fact{ + Kind: facts.KindRoute, + Name: "* " + m[1], + File: relFile, + Line: lineNum, + Props: map[string]any{ + "path": m[1], + "handler": m[2], + "framework": "django", + "language": "python", + }, + }) + } + } // Import: `import foo.bar` if m := importRe.FindStringSubmatch(line); m != nil { @@ -413,11 +598,85 @@ func extractFile(f *os.File, relFile string) []facts.Fact { } } + if err := scanner.Err(); err != nil { + log.Printf("[python-extractor] scanner error in %s: %v", relFile, err) + } + return result } // --- Helpers --- +// applyDecoratorProps sets structural boolean props on a symbol based on a +// decorator name. Only well-known structural decorators produce props; unknown +// decorators are silently ignored. +func applyDecoratorProps(props map[string]any, decoratorName string) { + // Use the last dot-separated component: "functools.cached_property" → "cached_property". + last := decoratorName + if idx := strings.LastIndex(decoratorName, "."); idx >= 0 { + last = decoratorName[idx+1:] + } + switch last { + case "property", "cached_property": + props["property"] = true + case "staticmethod": + props["static"] = true + case "classmethod": + props["class_method"] = true + case "abstractmethod": + props["abstract"] = true + case "task": + props["task"] = true + case "shared_task": + // shared_task is Celery-specific; bare @task is used by Airflow, Prefect, Luigi, etc. + props["task"] = true + props["framework"] = "celery" + } +} + +// detectDjango returns true if the project at repoPath uses Django, by scanning +// common dependency files and checking for manage.py. +func detectDjango(repoPath string) bool { + for _, name := range []string{"requirements.txt", "pyproject.toml", "setup.cfg", "setup.py"} { + data, err := os.ReadFile(filepath.Join(repoPath, name)) + if err != nil { + continue + } + if strings.Contains(strings.ToLower(string(data)), "django") { + return true + } + } + _, err := os.Stat(filepath.Join(repoPath, "manage.py")) + return err == nil +} + +// camelToSnake converts a PascalCase class name to the snake_case table name +// Django would auto-generate. e.g. "UserProfile" → "user_profile". +func camelToSnake(s string) string { + var b strings.Builder + for i := 0; i < len(s); i++ { + ch := s[i] + if i > 0 && ch >= 'A' && ch <= 'Z' { + b.WriteByte('_') + } + if ch >= 'A' && ch <= 'Z' { + b.WriteByte(ch + 32) // ASCII lowercase + } else { + b.WriteByte(ch) + } + } + return b.String() +} + +// lastComponent returns the last dot-separated segment of a qualified name. +// e.g. "models.Model" → "Model", "Model" → "Model". +func lastComponent(name string) string { + if idx := strings.LastIndex(name, "."); idx >= 0 { + return name[idx+1:] + } + return name +} + // buildQualName constructs a qualified name like "module.Outer.Inner.Name". // module is the file-based module path (e.g. "app/models/order" for "app/models/order.py"). func buildQualName(module string, stack []scopeEntry, name string) string { @@ -499,11 +758,11 @@ func splitBases(s string) []string { // stripGeneric removes generic type parameters from a base class name. // e.g. "Generic[T]" → "Generic", "CRUDBase[Model, Schema]" → "CRUDBase". func stripGeneric(s string) string { - if idx := strings.Index(s, "["); idx >= 0 { - return strings.TrimSpace(s[:idx]) + if before, _, ok := strings.Cut(s, "["); ok { + return strings.TrimSpace(before) } - if idx := strings.Index(s, "("); idx >= 0 { - return strings.TrimSpace(s[:idx]) + if before, _, ok := strings.Cut(s, "("); ok { + return strings.TrimSpace(before) } return strings.TrimSpace(s) } diff --git a/internal/extractors/pythonextractor/python_test.go b/internal/extractors/pythonextractor/python_test.go index 141202e..a26ed7a 100644 --- a/internal/extractors/pythonextractor/python_test.go +++ b/internal/extractors/pythonextractor/python_test.go @@ -81,7 +81,7 @@ class Order: defer f.Close() relFile := "app/models/order.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) // Class fact: module.ClassName @@ -127,7 +127,7 @@ async def fetch_data(url): defer f.Close() relFile := "services/utils.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) helperName := mod(relFile) + ".helper" @@ -162,7 +162,7 @@ class VespaSink(EmbeddingsSink): defer f.Close() relFile := "sinks/vespa_sink.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) clsName := mod(relFile) + ".VespaSink" @@ -187,7 +187,7 @@ class FeatureGroup(Base, TimestampMixin): defer f.Close() relFile := "db/models/feature_group.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) clsName := mod(relFile) + ".FeatureGroup" @@ -212,7 +212,7 @@ class CRUDEntity(CRUDBase[ModelType, IdType]): defer f.Close() relFile := "db/crud/crud.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) clsName := mod(relFile) + ".CRUDEntity" @@ -240,7 +240,7 @@ class Outer: defer f.Close() relFile := "pkg/nested.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) outerName := mod(relFile) + ".Outer" @@ -285,7 +285,7 @@ class Recommender: defer f.Close() relFile := "services/recommender.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) methodName := mod(relFile) + ".Recommender.recommend" @@ -311,7 +311,7 @@ import fastapi defer f.Close() relFile := "myapp/app.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) for _, target := range []string{"logging", "os", "fastapi"} { @@ -340,7 +340,7 @@ from .base import Base defer f.Close() relFile := "routes/routes.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) cases := []struct { @@ -381,7 +381,7 @@ async def health_check(): defer f.Close() relFile := "routes/health.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) routes := factsByKind(result, facts.KindRoute) if len(routes) != 1 { @@ -418,7 +418,7 @@ async def post_recommend_v2(body: RecommendV2Body) -> RecommendV2Response: defer f.Close() relFile := "routes/recommend.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) routes := factsByKind(result, facts.KindRoute) if len(routes) != 1 { @@ -452,7 +452,7 @@ async def delete_item(id: int): f := writeAndOpen(t, "items.py", src) defer f.Close() - result := extractFile(f, "routes/items.py") + result := extractFile(f, "routes/items.py", false) routes := factsByKind(result, facts.KindRoute) if len(routes) != 3 { @@ -481,7 +481,7 @@ async def login(): f := writeAndOpen(t, "auth.py", src) defer f.Close() - result := extractFile(f, "routes/auth.py") + result := extractFile(f, "routes/auth.py", false) routes := factsByKind(result, facts.KindRoute) if len(routes) != 1 { @@ -506,7 +506,7 @@ class FeatureGroup(Base): defer f.Close() relFile := "db/models/feature_group.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) storages := factsByKind(result, facts.KindStorage) if len(storages) != 1 { @@ -549,7 +549,7 @@ class MyService: defer f.Close() relFile := "services/service.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) // fake_def and FakeClass inside the docstring must NOT appear. @@ -583,7 +583,7 @@ class Validator: defer f.Close() relFile := "pkg/validator.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) validateName := mod(relFile) + ".Validator.validate" @@ -601,7 +601,7 @@ func TestExtractFile_LineNumbers(t *testing.T) { defer f.Close() relFile := "pkg/foo.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) clsName := mod(relFile) + ".Foo" @@ -634,7 +634,7 @@ class Foo: defer f.Close() relFile := "pkg/foo.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) clsName := mod(relFile) + ".Foo" @@ -678,7 +678,7 @@ async def post_recommend_v2( defer f.Close() relFile := "query_recommender/routes/recommend_v2.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) routes := factsByKind(result, facts.KindRoute) if len(routes) != 1 { @@ -731,7 +731,7 @@ class Entity(Base): defer f.Close() relFile := "db/models/entity.py" - result := extractFile(f, relFile) + result := extractFile(f, relFile, false) idx := byName(result) // Class with Base inheritance. @@ -917,3 +917,596 @@ func keys(m map[string]facts.Fact) []string { } return out } + +// --- Phase 1a: Decorator tracking --- + +func TestExtractFile_DecoratorProps_Property(t *testing.T) { + src := ` +class Config: + @property + def name(self) -> str: + return self._name + + @cached_property + def items(self): + return [] +` + f := writeAndOpen(t, "config.py", src) + defer f.Close() + + relFile := "app/config.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + for _, methodName := range []string{ + mod(relFile) + ".Config.name", + mod(relFile) + ".Config.items", + } { + m, ok := idx[methodName] + if !ok { + t.Fatalf("missing %q; keys: %v", methodName, keys(idx)) + } + if m.Props["property"] != true { + t.Errorf("%s: property = %v, want true", methodName, m.Props["property"]) + } + } +} + +func TestExtractFile_DecoratorProps_Staticmethod(t *testing.T) { + src := ` +class Utils: + @staticmethod + def parse(value): + return int(value) +` + f := writeAndOpen(t, "utils.py", src) + defer f.Close() + + relFile := "pkg/utils.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + methName := mod(relFile) + ".Utils.parse" + m, ok := idx[methName] + if !ok { + t.Fatalf("missing %q; keys: %v", methName, keys(idx)) + } + if m.Props["static"] != true { + t.Errorf("parse: static = %v, want true", m.Props["static"]) + } + if m.Props["class_method"] == true { + t.Error("parse: class_method should not be set") + } +} + +func TestExtractFile_DecoratorProps_Classmethod(t *testing.T) { + src := ` +class Repo: + @classmethod + def from_env(cls): + pass +` + f := writeAndOpen(t, "repo.py", src) + defer f.Close() + + relFile := "db/repo.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + methName := mod(relFile) + ".Repo.from_env" + m, ok := idx[methName] + if !ok { + t.Fatalf("missing %q; keys: %v", methName, keys(idx)) + } + if m.Props["class_method"] != true { + t.Errorf("from_env: class_method = %v, want true", m.Props["class_method"]) + } + if m.Props["static"] == true { + t.Error("from_env: static should not be set") + } +} + +func TestExtractFile_DecoratorProps_Abstractmethod(t *testing.T) { + src := ` +from abc import ABC, abstractmethod + +class Base(ABC): + @abstractmethod + def execute(self): + pass + + def concrete(self): + pass +` + f := writeAndOpen(t, "base.py", src) + defer f.Close() + + relFile := "core/base.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + executeName := mod(relFile) + ".Base.execute" + execute, ok := idx[executeName] + if !ok { + t.Fatalf("missing %q; keys: %v", executeName, keys(idx)) + } + if execute.Props["abstract"] != true { + t.Errorf("execute: abstract = %v, want true", execute.Props["abstract"]) + } + + // concrete() must NOT have abstract set. + concreteName := mod(relFile) + ".Base.concrete" + concrete, ok := idx[concreteName] + if !ok { + t.Fatalf("missing %q", concreteName) + } + if concrete.Props["abstract"] == true { + t.Error("concrete: abstract should not be set") + } +} + +func TestExtractFile_DecoratorProps_StackedDecorators(t *testing.T) { + // @classmethod + @abstractmethod on the same method — both props must be set. + src := ` +from abc import abstractmethod + +class Base: + @classmethod + @abstractmethod + def create(cls): + pass +` + f := writeAndOpen(t, "base.py", src) + defer f.Close() + + relFile := "core/base.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + methName := mod(relFile) + ".Base.create" + m, ok := idx[methName] + if !ok { + t.Fatalf("missing %q; keys: %v", methName, keys(idx)) + } + if m.Props["class_method"] != true { + t.Errorf("create: class_method = %v, want true", m.Props["class_method"]) + } + if m.Props["abstract"] != true { + t.Errorf("create: abstract = %v, want true", m.Props["abstract"]) + } +} + +func TestExtractFile_Task_Bare(t *testing.T) { + // Bare @task — framework-agnostic (Airflow, Prefect, etc.). + src := ` +@task +def process_records(): + pass +` + f := writeAndOpen(t, "tasks.py", src) + defer f.Close() + + relFile := "jobs/tasks.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + fnName := mod(relFile) + ".process_records" + fn, ok := idx[fnName] + if !ok { + t.Fatalf("missing %q; keys: %v", fnName, keys(idx)) + } + if fn.Props["task"] != true { + t.Errorf("process_records: task = %v, want true", fn.Props["task"]) + } + if fn.Props["framework"] != nil { + t.Errorf("process_records: framework = %v, want nil for bare @task", fn.Props["framework"]) + } +} + +func TestExtractFile_Task_SharedTask(t *testing.T) { + // @shared_task is Celery-specific and must set framework="celery". + src := ` +from celery import shared_task + +@shared_task +def send_welcome_email(user_id: int) -> None: + pass +` + f := writeAndOpen(t, "email_tasks.py", src) + defer f.Close() + + relFile := "notifications/email_tasks.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + fnName := mod(relFile) + ".send_welcome_email" + fn, ok := idx[fnName] + if !ok { + t.Fatalf("missing %q; keys: %v", fnName, keys(idx)) + } + if fn.Props["task"] != true { + t.Errorf("send_welcome_email: task = %v, want true", fn.Props["task"]) + } + if fn.Props["framework"] != "celery" { + t.Errorf("send_welcome_email: framework = %v, want celery", fn.Props["framework"]) + } +} + +func TestExtractFile_MultiLineDecorator(t *testing.T) { + // Multi-line decorator args must not clear pending state before the def. + // Without the bracket-depth fix the continuation lines clear pendingDecorators. + src := ` +@task( + bind=True, + max_retries=3, +) +def retry_task(self): + pass +` + f := writeAndOpen(t, "tasks.py", src) + defer f.Close() + + relFile := "jobs/tasks.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + fnName := mod(relFile) + ".retry_task" + fn, ok := idx[fnName] + if !ok { + t.Fatalf("missing %q; keys: %v", fnName, keys(idx)) + } + if fn.Props["task"] != true { + t.Errorf("retry_task: task = %v, want true (multi-line decorator must survive)", fn.Props["task"]) + } +} + +// --- Phase 1b: Return type hints --- + +func TestExtractFile_ReturnType_SingleLine(t *testing.T) { + src := ` +def is_ready() -> bool: + return True + +def get_count() -> int: + return 42 + +def no_annotation(): + pass +` + f := writeAndOpen(t, "funcs.py", src) + defer f.Close() + + relFile := "pkg/funcs.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + cases := []struct { + name string + want string + }{ + {mod(relFile) + ".is_ready", "bool"}, + {mod(relFile) + ".get_count", "int"}, + } + for _, tc := range cases { + fn, ok := idx[tc.name] + if !ok { + t.Fatalf("missing %q", tc.name) + } + if fn.Props["return_type"] != tc.want { + t.Errorf("%s: return_type = %v, want %q", tc.name, fn.Props["return_type"], tc.want) + } + } + + // no_annotation must have no return_type prop. + noAnn := idx[mod(relFile)+".no_annotation"] + if noAnn.Props["return_type"] != nil { + t.Errorf("no_annotation: return_type = %v, want nil", noAnn.Props["return_type"]) + } +} + +func TestExtractFile_ReturnType_Complex(t *testing.T) { + src := ` +def get_config() -> dict[str, Any]: + pass + +def find_user() -> Optional[str]: + pass + +def get_items() -> list[str] | None: + pass +` + f := writeAndOpen(t, "service.py", src) + defer f.Close() + + relFile := "svc/service.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + cases := []struct{ name, want string }{ + {mod(relFile) + ".get_config", "dict[str, Any]"}, + {mod(relFile) + ".find_user", "Optional[str]"}, + {mod(relFile) + ".get_items", "list[str] | None"}, + } + for _, tc := range cases { + fn, ok := idx[tc.name] + if !ok { + t.Fatalf("missing %q", tc.name) + } + if fn.Props["return_type"] != tc.want { + t.Errorf("%s: return_type = %v, want %q", tc.name, fn.Props["return_type"], tc.want) + } + } +} + +func TestExtractFile_ReturnType_MultiLine(t *testing.T) { + // Return type on the closing paren line of a multi-line signature. + src := ` +def create_handler( + request: Request, + response: Response, +) -> Optional[str]: + pass +` + f := writeAndOpen(t, "handler.py", src) + defer f.Close() + + relFile := "api/handler.py" + result := extractFile(f, relFile, false) + idx := byName(result) + + fnName := mod(relFile) + ".create_handler" + fn, ok := idx[fnName] + if !ok { + t.Fatalf("missing %q; keys: %v", fnName, keys(idx)) + } + if fn.Props["return_type"] != "Optional[str]" { + t.Errorf("create_handler: return_type = %v, want Optional[str]", fn.Props["return_type"]) + } +} + +// --- Phase 1c: Django support --- + +func TestExtractFile_DjangoModel(t *testing.T) { + src := ` +from django.db import models + +class Order(models.Model): + total = models.DecimalField(max_digits=10, decimal_places=2) + +class UserProfile(models.Model): + user = models.OneToOneField('User', on_delete=models.CASCADE) +` + f := writeAndOpen(t, "models.py", src) + defer f.Close() + + relFile := "shop/models.py" + result := extractFile(f, relFile, true) + + storages := factsByKind(result, facts.KindStorage) + if len(storages) != 2 { + t.Fatalf("expected 2 storage facts, got %d: %v", len(storages), storages) + } + + idx := byName(result) + + // Order → "order" (camelToSnake) + order, ok := idx["order"] + if !ok { + t.Fatalf("missing storage fact %q; keys: %v", "order", keys(idx)) + } + if order.Props["framework"] != "django" { + t.Errorf("order: framework = %v, want django", order.Props["framework"]) + } + if order.Props["storage_kind"] != "table" { + t.Errorf("order: storage_kind = %v, want table", order.Props["storage_kind"]) + } + wantClass := mod(relFile) + ".Order" + if order.Props["class"] != wantClass { + t.Errorf("order: class = %v, want %q", order.Props["class"], wantClass) + } + + // UserProfile → "user_profile" + if _, ok := idx["user_profile"]; !ok { + t.Errorf("missing storage fact %q; keys: %v", "user_profile", keys(idx)) + } +} + +func TestExtractFile_DjangoCBV(t *testing.T) { + src := ` +from rest_framework.views import APIView + +class OrderView(APIView): + def get(self, request): + pass +` + f := writeAndOpen(t, "views.py", src) + defer f.Close() + + relFile := "shop/views.py" + result := extractFile(f, relFile, true) + idx := byName(result) + + clsName := mod(relFile) + ".OrderView" + cls, ok := idx[clsName] + if !ok { + t.Fatalf("missing %q; keys: %v", clsName, keys(idx)) + } + if cls.Props["django_component"] != "view" { + t.Errorf("OrderView: django_component = %v, want view", cls.Props["django_component"]) + } + if cls.Props["framework"] != "django" { + t.Errorf("OrderView: framework = %v, want django", cls.Props["framework"]) + } +} + +func TestExtractFile_DRFSerializer(t *testing.T) { + src := ` +from rest_framework import serializers + +class OrderSerializer(serializers.ModelSerializer): + class Meta: + model = Order + fields = '__all__' +` + f := writeAndOpen(t, "serializers.py", src) + defer f.Close() + + relFile := "shop/serializers.py" + result := extractFile(f, relFile, true) + idx := byName(result) + + clsName := mod(relFile) + ".OrderSerializer" + cls, ok := idx[clsName] + if !ok { + t.Fatalf("missing %q; keys: %v", clsName, keys(idx)) + } + if cls.Props["django_component"] != "serializer" { + t.Errorf("OrderSerializer: django_component = %v, want serializer", cls.Props["django_component"]) + } + if cls.Props["framework"] != "django" { + t.Errorf("OrderSerializer: framework = %v, want django", cls.Props["framework"]) + } +} + +func TestExtractFile_DjangoURL(t *testing.T) { + src := ` +from django.urls import path +from . import views + +urlpatterns = [ + path('orders/', views.OrderListView.as_view()), + path('orders//', views.OrderDetailView.as_view()), + re_path(r'^legacy/$', views.legacy_view), +] +` + // File must be named urls.py for Django URL extraction. + f := writeAndOpen(t, "urls.py", src) + defer f.Close() + + result := extractFile(f, "shop/urls.py", true) + routes := factsByKind(result, facts.KindRoute) + + if len(routes) != 3 { + t.Fatalf("expected 3 route facts, got %d: %v", len(routes), routes) + } + idx := byName(result) + + for _, wantName := range []string{"* orders/", "* orders//", "* ^legacy/$"} { + r, ok := idx[wantName] + if !ok { + t.Errorf("missing route %q; keys: %v", wantName, keys(idx)) + continue + } + if r.Props["framework"] != "django" { + t.Errorf("%s: framework = %v, want django", wantName, r.Props["framework"]) + } + } +} + +func TestExtractFile_DjangoURL_NonURLsFile(t *testing.T) { + // Django URL patterns in a file not named urls.py must NOT produce route facts. + src := ` +urlpatterns = [ + path('orders/', views.OrderListView.as_view()), +] +` + f := writeAndOpen(t, "routing.py", src) + defer f.Close() + + result := extractFile(f, "shop/routing.py", true) + routes := factsByKind(result, facts.KindRoute) + if len(routes) != 0 { + t.Errorf("expected no routes in non-urls.py file, got %d", len(routes)) + } +} + +func TestExtractFile_DjangoAPIView(t *testing.T) { + src := ` +from rest_framework.decorators import api_view + +@api_view(['GET', 'POST']) +def order_list(request): + pass + +@api_view(['GET']) +def order_detail(request, pk): + pass +` + f := writeAndOpen(t, "views.py", src) + defer f.Close() + + relFile := "shop/views.py" + result := extractFile(f, relFile, true) + routes := factsByKind(result, facts.KindRoute) + + // order_list has GET+POST → 2 routes; order_detail has GET → 1 route. + if len(routes) != 3 { + t.Fatalf("expected 3 route facts, got %d: %v", len(routes), routes) + } + + handlerBase := mod(relFile) + ".order_list" + idx := byName(result) + for _, wantName := range []string{ + "GET (view) " + handlerBase, + "POST (view) " + handlerBase, + } { + r, ok := idx[wantName] + if !ok { + t.Errorf("missing route %q; keys: %v", wantName, keys(idx)) + continue + } + if r.Props["framework"] != "django" { + t.Errorf("%s: framework = %v, want django", wantName, r.Props["framework"]) + } + } +} + +// --- Helper unit tests --- + +func TestCamelToSnake(t *testing.T) { + cases := []struct{ input, want string }{ + {"Order", "order"}, + {"UserProfile", "user_profile"}, + {"ProductCategory", "product_category"}, + {"Foo", "foo"}, + {"FooBar", "foo_bar"}, + } + for _, tc := range cases { + got := camelToSnake(tc.input) + if got != tc.want { + t.Errorf("camelToSnake(%q) = %q, want %q", tc.input, got, tc.want) + } + } +} + +func TestDetectDjango(t *testing.T) { + t.Run("requirements_txt", func(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "requirements.txt"), []byte("django>=4.2\nrest_framework\n"), 0o644); err != nil { + t.Fatal(err) + } + if !detectDjango(dir) { + t.Error("detectDjango should return true for requirements.txt with django") + } + }) + + t.Run("manage_py", func(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "manage.py"), []byte("#!/usr/bin/env python\n"), 0o644); err != nil { + t.Fatal(err) + } + if !detectDjango(dir) { + t.Error("detectDjango should return true when manage.py is present") + } + }) + + t.Run("no_django", func(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "requirements.txt"), []byte("fastapi\nsqlalchemy\n"), 0o644); err != nil { + t.Fatal(err) + } + if detectDjango(dir) { + t.Error("detectDjango should return false for non-Django project") + } + }) +} From 836496a77295591507de653f1e1250d48d362152 Mon Sep 17 00:00:00 2001 From: GertL Date: Sun, 14 Jun 2026 18:12:04 +0200 Subject: [PATCH 2/5] Add Python AST extraction and related tests for improved fact generation --- go.mod | 1 + internal/extractors/pythonextractor/python.go | 17 +- .../extractors/pythonextractor/python_ast.go | 653 ++++++++++++++++++ .../pythonextractor/python_ast_test.go | 301 ++++++++ .../extractors/pythonextractor/python_test.go | 16 - 5 files changed, 971 insertions(+), 17 deletions(-) create mode 100644 internal/extractors/pythonextractor/python_ast.go create mode 100644 internal/extractors/pythonextractor/python_ast_test.go diff --git a/go.mod b/go.mod index 2d56200..d8266ba 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( github.com/modelcontextprotocol/go-sdk v1.3.0 github.com/tree-sitter-grammars/tree-sitter-kotlin v1.1.0 github.com/tree-sitter/go-tree-sitter v0.24.0 + github.com/tree-sitter/tree-sitter-python v0.21.1-0.20240818005537-55a9b8a4fbfb github.com/tree-sitter/tree-sitter-typescript v0.23.2 gopkg.in/yaml.v3 v3.0.1 ) diff --git a/internal/extractors/pythonextractor/python.go b/internal/extractors/pythonextractor/python.go index ae17526..cb5ef24 100644 --- a/internal/extractors/pythonextractor/python.go +++ b/internal/extractors/pythonextractor/python.go @@ -3,6 +3,7 @@ package pythonextractor import ( "bufio" "context" + "io" "io/fs" "log" "os" @@ -93,8 +94,14 @@ func (e *PythonExtractor) Extract(ctx context.Context, repoPath string, files [] continue } - fileFacts := extractFile(f, relFile, isDjango) + src, readErr := readAll(f) f.Close() + var fileFacts []facts.Fact + if readErr != nil { + log.Printf("[python-extractor] error reading %s: %v", relFile, readErr) + continue + } + fileFacts = extractFileAST(src, relFile, isDjango) allFacts = append(allFacts, fileFacts...) dir := filepath.Dir(relFile) @@ -771,3 +778,11 @@ func stripGeneric(s string) string { func isPythonFile(path string) bool { return strings.HasSuffix(strings.ToLower(path), ".py") } + +// readAll reads all bytes from an open file, seeking to the start first. +func readAll(f *os.File) ([]byte, error) { + if _, err := f.Seek(0, io.SeekStart); err != nil { + return nil, err + } + return io.ReadAll(f) +} diff --git a/internal/extractors/pythonextractor/python_ast.go b/internal/extractors/pythonextractor/python_ast.go new file mode 100644 index 0000000..dede3fc --- /dev/null +++ b/internal/extractors/pythonextractor/python_ast.go @@ -0,0 +1,653 @@ +package pythonextractor + +import ( + "path/filepath" + "strings" + "unicode" + + "github.com/enola-labs/enola/internal/facts" + python "github.com/tree-sitter/tree-sitter-python/bindings/go" + sitter "github.com/tree-sitter/go-tree-sitter" +) + +// extractFileAST parses a Python file with tree-sitter and emits architectural +// facts. It is a superset of extractFile: every symbol / import / route / storage +// fact is preserved, and RelCalls / RelInstantiates edges are added when call +// sites are observed inside function bodies. +func extractFileAST(src []byte, relFile string, isDjango bool) []facts.Fact { + parser := sitter.NewParser() + defer parser.Close() + if err := parser.SetLanguage(sitter.NewLanguage(python.Language())); err != nil { + return nil + } + + tree := parser.Parse(src, nil) + defer tree.Close() + + module := strings.TrimSuffix(relFile, ".py") + dir := filepath.Dir(relFile) + + w := &pyWalker{ + src: src, + relFile: relFile, + module: module, + dir: dir, + isDjango: isDjango, + } + w.walkModule(tree.RootNode()) + return w.out +} + +type pyWalker struct { + src []byte + relFile string + module string + dir string + isDjango bool + + out []facts.Fact + + // typeStack holds enclosing class names so methods get qualified names. + typeStack []string + + // ownerStack: top element is the fact that receives RelCalls / RelInstantiates + // discovered while walking its body. + ownerStack []*facts.Fact + + // importMap maps a local name to its canonical fact target (empty = external). + importMap map[string]string + + // methodSets[i] is the set of methods declared directly in typeStack[i], + // used to resolve bare same-class calls. + methodSets []map[string]bool +} + +func (w *pyWalker) pushOwner(f *facts.Fact) { w.ownerStack = append(w.ownerStack, f) } +func (w *pyWalker) popOwner() { w.ownerStack = w.ownerStack[:len(w.ownerStack)-1] } +func (w *pyWalker) currentOwner() *facts.Fact { + if len(w.ownerStack) == 0 { + return nil + } + return w.ownerStack[len(w.ownerStack)-1] +} + +func (w *pyWalker) enclosingType() string { return strings.Join(w.typeStack, ".") } + +func (w *pyWalker) qualify(name string) string { + if t := w.enclosingType(); t != "" { + return t + "." + name + } + return name +} + +func (w *pyWalker) pushType(name string, methods map[string]bool) { + w.typeStack = append(w.typeStack, name) + w.methodSets = append(w.methodSets, methods) +} + +func (w *pyWalker) popType() { + w.typeStack = w.typeStack[:len(w.typeStack)-1] + w.methodSets = w.methodSets[:len(w.methodSets)-1] +} + +func (w *pyWalker) currentMethods() map[string]bool { + if len(w.methodSets) == 0 { + return nil + } + return w.methodSets[len(w.methodSets)-1] +} + +// walkModule iterates the top-level statements of a module node. +func (w *pyWalker) walkModule(root *sitter.Node) { + for i := uint(0); i < uint(root.ChildCount()); i++ { + w.walkStatement(root.Child(i)) + } +} + +func (w *pyWalker) walkStatement(node *sitter.Node) { + if node == nil { + return + } + switch node.Kind() { + case "import_statement": + w.handleImport(node) + case "import_from_statement": + w.handleFromImport(node) + case "class_definition": + w.handleClass(node, nil) + case "function_definition": + w.handleFunction(node, nil) + case "decorated_definition": + w.handleDecoratedDefinition(node) + case "expression_statement": + // __tablename__ = "foo" (SQLAlchemy) lives here at class body level. + w.handleExprStatement(node) + case "block": + for i := uint(0); i < uint(node.ChildCount()); i++ { + w.walkStatement(node.Child(i)) + } + } +} + +// handleImport handles `import foo.bar` — emits KindDependency + RelImports. +func (w *pyWalker) handleImport(node *sitter.Node) { + for i := uint(0); i < uint(node.ChildCount()); i++ { + c := node.Child(i) + if c.Kind() == "dotted_name" || c.Kind() == "aliased_import" { + var name, alias string + if c.Kind() == "aliased_import" { + nameNode := c.ChildByFieldName("name") + aliasNode := c.ChildByFieldName("alias") + if nameNode == nil { + continue + } + name = pyText(c.ChildByFieldName("name"), w.src) + if aliasNode != nil { + alias = pyText(aliasNode, w.src) + } + } else { + name = pyText(c, w.src) + } + target := w.dir + " -> " + name + w.out = append(w.out, facts.Fact{ + Kind: facts.KindDependency, + Name: target, + File: w.relFile, + Line: int(node.StartPosition().Row) + 1, + Props: map[string]any{"language": "python"}, + Relations: []facts.Relation{ + {Kind: facts.RelImports, Target: name}, + }, + }) + local := alias + if local == "" { + if dot := strings.LastIndex(name, "."); dot >= 0 { + local = name[dot+1:] + } else { + local = name + } + } + w.setImport(local, "") + } + } +} + +// handleFromImport handles `from foo.bar import Baz, Qux`. +func (w *pyWalker) handleFromImport(node *sitter.Node) { + moduleNode := node.ChildByFieldName("module_name") + if moduleNode == nil { + return + } + moduleName := pyText(moduleNode, w.src) + + // Determine if this is an intra-project import (relative or same-tree dotted). + isRelative := strings.HasPrefix(moduleName, ".") || + strings.HasPrefix(pyText(node, w.src), "from .") + + target := w.dir + " -> " + moduleName + w.out = append(w.out, facts.Fact{ + Kind: facts.KindDependency, + Name: target, + File: w.relFile, + Line: int(node.StartPosition().Row) + 1, + Props: map[string]any{"language": "python"}, + Relations: []facts.Relation{ + {Kind: facts.RelImports, Target: moduleName}, + }, + }) + + // Map each imported name to a resolvable target or "" (external). + for i := uint(0); i < uint(node.ChildCount()); i++ { + c := node.Child(i) + if c.Kind() != "dotted_name" && c.Kind() != "identifier" && c.Kind() != "aliased_import" { + continue + } + var localName, importedName string + if c.Kind() == "aliased_import" { + n := c.ChildByFieldName("name") + a := c.ChildByFieldName("alias") + if n == nil { + continue + } + importedName = pyText(n, w.src) + if a != nil { + localName = pyText(a, w.src) + } else { + localName = importedName + } + } else { + importedName = pyText(c, w.src) + localName = importedName + } + + if isRelative { + // Relative import → resolve to a local module path. + base := moduleName + if strings.HasPrefix(base, ".") { + base = w.dir + "/" + strings.TrimLeft(base, ".") + } + w.setImport(localName, base+"."+importedName) + } else { + // External or ambiguous — suppress call edges to this name. + w.setImport(localName, "") + } + } +} + +func (w *pyWalker) setImport(local, target string) { + if local == "" || local == "*" { + return + } + if w.importMap == nil { + w.importMap = make(map[string]string) + } + w.importMap[local] = target +} + +// handleDecoratedDefinition unwraps `@decorator\ndef/class ...` nodes. +func (w *pyWalker) handleDecoratedDefinition(node *sitter.Node) { + var decorators []string + var pendingApiViewMethods []string + + for i := uint(0); i < uint(node.ChildCount()); i++ { + c := node.Child(i) + switch c.Kind() { + case "decorator": + text := pyText(c, w.src) + // FastAPI / Starlette route decorator. + if m := routeDecoratorRe.FindStringSubmatch(text); m != nil { + w.out = append(w.out, facts.Fact{ + Kind: facts.KindRoute, + Name: strings.ToUpper(m[2]) + " " + m[3], + File: w.relFile, + Line: int(c.StartPosition().Row) + 1, + Props: map[string]any{ + "method": strings.ToUpper(m[2]), + "path": m[3], + "framework": "fastapi", + }, + }) + continue + } + // DRF @api_view(['GET','POST']). + if m := apiViewRe.FindStringSubmatch(text); m != nil { + for _, meth := range httpMethodWordRe.FindAllString(m[1], -1) { + pendingApiViewMethods = append(pendingApiViewMethods, meth) + } + continue + } + // Generic decorator name capture. + if m := decoratorRe.FindStringSubmatch(text); m != nil { + decorators = append(decorators, m[1]) + } + + case "function_definition": + w.handleFunction(c, decorators) + // @api_view routes — emit after we know the handler name. + if len(pendingApiViewMethods) > 0 { + handlerName := w.module + "." + w.qualify(pyFuncName(c, w.src)) + for _, meth := range pendingApiViewMethods { + w.out = append(w.out, facts.Fact{ + Kind: facts.KindRoute, + Name: meth + " (view) " + handlerName, + File: w.relFile, + Line: int(c.StartPosition().Row) + 1, + Props: map[string]any{ + "method": meth, + "framework": "django", + "handler": handlerName, + }, + }) + } + } + + case "class_definition": + w.handleClass(c, decorators) + } + } +} + +// handleClass emits a KindSymbol fact for a class and walks its body. +func (w *pyWalker) handleClass(node *sitter.Node, decorators []string) { + nameNode := node.ChildByFieldName("name") + if nameNode == nil { + return + } + name := pyText(nameNode, w.src) + qualName := w.module + "." + w.qualify(name) + + props := map[string]any{ + "symbol_kind": facts.SymbolClass, + "language": "python", + } + rels := []facts.Relation{{Kind: facts.RelDeclares, Target: w.dir}} + + // Superclasses. + var bases []string + if args := node.ChildByFieldName("superclasses"); args != nil { + for i := uint(0); i < uint(args.ChildCount()); i++ { + c := args.Child(i) + if c.Kind() == "identifier" || c.Kind() == "attribute" { + base := pyText(c, w.src) + bases = append(bases, base) + rels = append(rels, facts.Relation{Kind: facts.RelImplements, Target: base}) + } + } + } + + for _, dec := range decorators { + applyDecoratorProps(props, dec) + } + + // Django classification. + if w.isDjango { + for _, base := range bases { + last := lastComponent(base) + if djangoModelBases[last] { + props["framework"] = "django" + tableName := camelToSnake(name) + w.out = append(w.out, facts.Fact{ + Kind: facts.KindStorage, + Name: tableName, + File: w.relFile, + Line: int(node.StartPosition().Row) + 1, + Props: map[string]any{ + "storage_kind": "table", + "framework": "django", + "class": qualName, + }, + }) + break + } + if djangoCBVBases[last] { + props["django_component"] = "view" + props["framework"] = "django" + break + } + if djangoSerializerBases[last] { + props["django_component"] = "serializer" + props["framework"] = "django" + break + } + } + } + + // Django urls.py: emit route facts from path()/re_path() calls in the class body. + if w.isDjango && filepath.Base(w.relFile) == "urls.py" { + bodyNode := node.ChildByFieldName("body") + if bodyNode != nil { + bodyText := pyText(bodyNode, w.src) + for _, m := range urlPathRe.FindAllStringSubmatch(bodyText, -1) { + w.out = append(w.out, facts.Fact{ + Kind: facts.KindRoute, + Name: "* " + m[1], + File: w.relFile, + Line: int(node.StartPosition().Row) + 1, + Props: map[string]any{ + "path": m[1], + "handler": m[2], + "framework": "django", + }, + }) + } + } + } + + f := facts.Fact{ + Kind: facts.KindSymbol, + Name: qualName, + File: w.relFile, + Line: int(node.StartPosition().Row) + 1, + Props: props, + Relations: rels, + } + + w.out = append(w.out, f) + owner := &w.out[len(w.out)-1] + w.pushOwner(owner) + + bodyNode := node.ChildByFieldName("body") + w.pushType(name, collectPyMethodNames(bodyNode, w.src)) + if bodyNode != nil { + w.walkBody(bodyNode) + } + w.popType() + w.popOwner() +} + +// handleFunction emits a KindSymbol fact for a function/method. +func (w *pyWalker) handleFunction(node *sitter.Node, decorators []string) { + nameNode := node.ChildByFieldName("name") + if nameNode == nil { + return + } + name := pyText(nameNode, w.src) + qualName := w.module + "." + w.qualify(name) + + props := map[string]any{ + "symbol_kind": facts.SymbolFunc, + "language": "python", + } + if len(w.typeStack) > 0 { + props["receiver"] = w.typeStack[len(w.typeStack)-1] + } + + // async keyword: look for it as a sibling before the `def` keyword. + fullText := pyText(node, w.src) + if strings.HasPrefix(strings.TrimSpace(fullText), "async ") { + props["async"] = true + } + + // Return type. + if retNode := node.ChildByFieldName("return_type"); retNode != nil { + rt := strings.TrimSpace(pyText(retNode, w.src)) + if strings.HasPrefix(rt, "->") { + rt = strings.TrimSpace(rt[2:]) + } + if rt != "" { + props["return_type"] = rt + } + } + + for _, dec := range decorators { + applyDecoratorProps(props, dec) + } + + rels := []facts.Relation{{Kind: facts.RelDeclares, Target: w.dir}} + + f := facts.Fact{ + Kind: facts.KindSymbol, + Name: qualName, + File: w.relFile, + Line: int(node.StartPosition().Row) + 1, + Props: props, + Relations: rels, + } + + w.out = append(w.out, f) + owner := &w.out[len(w.out)-1] + w.pushOwner(owner) + if bodyNode := node.ChildByFieldName("body"); bodyNode != nil { + w.walkForCalls(bodyNode) + } + w.popOwner() +} + +// handleExprStatement checks for SQLAlchemy __tablename__ assignments. +func (w *pyWalker) handleExprStatement(node *sitter.Node) { + text := pyText(node, w.src) + if m := tableNameRe.FindStringSubmatch(text); m != nil { + w.out = append(w.out, facts.Fact{ + Kind: facts.KindStorage, + Name: m[1], + File: w.relFile, + Line: int(node.StartPosition().Row) + 1, + Props: map[string]any{ + "storage_kind": "table", + "framework": "sqlalchemy", + }, + }) + } +} + +// walkBody walks a class body, dispatching each statement. +func (w *pyWalker) walkBody(body *sitter.Node) { + for i := uint(0); i < uint(body.ChildCount()); i++ { + w.walkStatement(body.Child(i)) + } +} + +// walkForCalls recursively scans a function body for call nodes and emits +// RelCalls / RelInstantiates on the current owner. +func (w *pyWalker) walkForCalls(node *sitter.Node) { + if node == nil { + return + } + if node.Kind() == "call" { + if fn := node.ChildByFieldName("function"); fn != nil { + w.emitCallEdge(fn) + } + } + // Don't recurse into nested class/function definitions — they get their own owner. + switch node.Kind() { + case "class_definition", "function_definition", "decorated_definition": + return + } + for i := uint(0); i < uint(node.ChildCount()); i++ { + w.walkForCalls(node.Child(i)) + } +} + +// emitCallEdge resolves the callee node and appends a relation to the current owner. +func (w *pyWalker) emitCallEdge(fn *sitter.Node) { + owner := w.currentOwner() + if owner == nil { + return + } + + switch fn.Kind() { + case "identifier": + name := pyText(fn, w.src) + if pyBuiltins[name] { + return + } + if pyCapitalized(name) { + owner.Relations = append(owner.Relations, facts.Relation{ + Kind: facts.RelInstantiates, + Target: name, + }) + return + } + if target := w.resolveCall(name); target != "" { + owner.Relations = append(owner.Relations, facts.Relation{ + Kind: facts.RelCalls, + Target: target, + }) + } + + case "attribute": + // self.method() or obj.method() — only resolve self.method. + objNode := fn.ChildByFieldName("object") + attrNode := fn.ChildByFieldName("attribute") + if objNode == nil || attrNode == nil { + return + } + obj := pyText(objNode, w.src) + attr := pyText(attrNode, w.src) + if obj == "self" || obj == "cls" { + if methods := w.currentMethods(); methods[attr] { + owner.Relations = append(owner.Relations, facts.Relation{ + Kind: facts.RelCalls, + Target: w.module + "." + w.enclosingType() + "." + attr, + }) + } + } + } +} + +// resolveCall maps a bare call name to a canonical fact target. +func (w *pyWalker) resolveCall(name string) string { + // Same-class method. + if methods := w.currentMethods(); methods[name] { + return w.module + "." + w.enclosingType() + "." + name + } + // Imported name. + if target, ok := w.importMap[name]; ok { + return target // "" means external → no edge + } + // Same-module top-level function. + return w.module + "." + name +} + +// collectPyMethodNames returns the set of function names declared directly in a +// class body node. +func collectPyMethodNames(body *sitter.Node, src []byte) map[string]bool { + methods := make(map[string]bool) + if body == nil { + return methods + } + for i := uint(0); i < uint(body.ChildCount()); i++ { + c := body.Child(i) + var fn *sitter.Node + switch c.Kind() { + case "function_definition": + fn = c + case "decorated_definition": + for j := uint(0); j < uint(c.ChildCount()); j++ { + if c.Child(j).Kind() == "function_definition" { + fn = c.Child(j) + break + } + } + } + if fn != nil { + if nameNode := fn.ChildByFieldName("name"); nameNode != nil { + methods[pyText(nameNode, src)] = true + } + } + } + return methods +} + +func pyFuncName(node *sitter.Node, src []byte) string { + if n := node.ChildByFieldName("name"); n != nil { + return pyText(n, src) + } + return "" +} + +func pyText(node *sitter.Node, src []byte) string { + if node == nil { + return "" + } + return string(src[node.StartByte():node.EndByte()]) +} + +func pyCapitalized(s string) bool { + if s == "" { + return false + } + return unicode.IsUpper([]rune(s)[0]) +} + +// pyBuiltins are Python built-in functions that appear as bare calls without +// an import and have no local fact — resolving them would produce phantom edges. +var pyBuiltins = map[string]bool{ + "print": true, "len": true, "range": true, "enumerate": true, "zip": true, + "map": true, "filter": true, "sorted": true, "reversed": true, "list": true, + "dict": true, "set": true, "tuple": true, "str": true, "int": true, + "float": true, "bool": true, "bytes": true, "type": true, "isinstance": true, + "issubclass": true, "hasattr": true, "getattr": true, "setattr": true, + "delattr": true, "callable": true, "repr": true, "hash": true, "id": true, + "abs": true, "round": true, "min": true, "max": true, "sum": true, + "any": true, "all": true, "next": true, "iter": true, "open": true, + "super": true, "object": true, "property": true, "staticmethod": true, + "classmethod": true, "vars": true, "dir": true, "globals": true, + "locals": true, "exec": true, "eval": true, "compile": true, + "input": true, "format": true, "chr": true, "ord": true, "hex": true, + "oct": true, "bin": true, "pow": true, "divmod": true, "slice": true, + "NotImplemented": true, "Exception": true, "ValueError": true, + "TypeError": true, "KeyError": true, "IndexError": true, + "AttributeError": true, "RuntimeError": true, "StopIteration": true, + "GeneratorExit": true, "SystemExit": true, "KeyboardInterrupt": true, +} diff --git a/internal/extractors/pythonextractor/python_ast_test.go b/internal/extractors/pythonextractor/python_ast_test.go new file mode 100644 index 0000000..e68e9e0 --- /dev/null +++ b/internal/extractors/pythonextractor/python_ast_test.go @@ -0,0 +1,301 @@ +package pythonextractor + +import ( + "os" + "path/filepath" + "testing" + + "github.com/enola-labs/enola/internal/facts" +) + +// astExtract is a helper that writes src to a temp file and runs extractFileAST. +func astExtract(t *testing.T, filename, src string, isDjango bool) []facts.Fact { + t.Helper() + dir := t.TempDir() + path := filepath.Join(dir, filename) + if err := os.WriteFile(path, []byte(src), 0o644); err != nil { + t.Fatal(err) + } + return extractFileAST([]byte(src), filename, isDjango) +} + +// relsByKind returns all relations of a given kind from a fact. +func relsByKind(f facts.Fact, kind string) []string { + var out []string + for _, r := range f.Relations { + if r.Kind == kind { + out = append(out, r.Target) + } + } + return out +} + +// --- Call graph tests --- + +func TestAST_SameModuleFunctionCall(t *testing.T) { + src := ` +def helper(): + pass + +def main(): + helper() +` + result := astExtract(t, "svc.py", src, false) + idx := byName(result) + + mainFact, ok := idx["svc.main"] + if !ok { + t.Fatalf("missing svc.main; keys: %v", keys(idx)) + } + calls := relsByKind(mainFact, facts.RelCalls) + if len(calls) == 0 { + t.Fatal("svc.main: expected RelCalls to svc.helper, got none") + } + found := false + for _, c := range calls { + if c == "svc.helper" { + found = true + } + } + if !found { + t.Errorf("svc.main: RelCalls = %v, want svc.helper", calls) + } +} + +func TestAST_SelfMethodCall(t *testing.T) { + src := ` +class Service: + def _do_work(self): + pass + + def run(self): + self._do_work() +` + result := astExtract(t, "svc.py", src, false) + idx := byName(result) + + runFact, ok := idx["svc.Service.run"] + if !ok { + t.Fatalf("missing svc.Service.run; keys: %v", keys(idx)) + } + calls := relsByKind(runFact, facts.RelCalls) + found := false + for _, c := range calls { + if c == "svc.Service._do_work" { + found = true + } + } + if !found { + t.Errorf("Service.run: RelCalls = %v, want svc.Service._do_work", calls) + } +} + +func TestAST_Constructor_RelInstantiates(t *testing.T) { + src := ` +class Order: + pass + +def create(): + o = Order() + return o +` + result := astExtract(t, "models.py", src, false) + idx := byName(result) + + createFact, ok := idx["models.create"] + if !ok { + t.Fatalf("missing models.create; keys: %v", keys(idx)) + } + insts := relsByKind(createFact, facts.RelInstantiates) + found := false + for _, i := range insts { + if i == "Order" { + found = true + } + } + if !found { + t.Errorf("create: RelInstantiates = %v, want Order", insts) + } +} + +func TestAST_NoEdgeForBuiltins(t *testing.T) { + src := ` +def process(items): + result = list(map(str, items)) + print(len(result)) + return sorted(result) +` + result := astExtract(t, "util.py", src, false) + idx := byName(result) + + fn, ok := idx["util.process"] + if !ok { + t.Fatalf("missing util.process") + } + calls := relsByKind(fn, facts.RelCalls) + for _, c := range calls { + if c == "util.list" || c == "util.print" || c == "util.sorted" || c == "util.map" || c == "util.str" || c == "util.len" { + t.Errorf("process: should not emit call edge to builtin, got %q", c) + } + } +} + +func TestAST_ReturnType_FromAST(t *testing.T) { + // tree-sitter reads the return type node directly — no regex needed. + src := ` +def get_user(user_id: int) -> Optional[str]: + pass + +def create_order( + items: list, + total: float, +) -> dict[str, Any]: + pass +` + result := astExtract(t, "api.py", src, false) + idx := byName(result) + + cases := []struct{ name, want string }{ + {"api.get_user", "Optional[str]"}, + {"api.create_order", "dict[str, Any]"}, + } + for _, tc := range cases { + fn, ok := idx[tc.name] + if !ok { + t.Fatalf("missing %q; keys: %v", tc.name, keys(idx)) + } + if fn.Props["return_type"] != tc.want { + t.Errorf("%s: return_type = %v, want %q", tc.name, fn.Props["return_type"], tc.want) + } + } +} + +func TestAST_NestedClass(t *testing.T) { + src := ` +class Outer: + class Inner: + def method(self): + pass +` + result := astExtract(t, "nested.py", src, false) + idx := byName(result) + + if _, ok := idx["nested.Outer"]; !ok { + t.Errorf("missing nested.Outer; keys: %v", keys(idx)) + } + if _, ok := idx["nested.Outer.Inner"]; !ok { + t.Errorf("missing nested.Outer.Inner; keys: %v", keys(idx)) + } + if _, ok := idx["nested.Outer.Inner.method"]; !ok { + t.Errorf("missing nested.Outer.Inner.method; keys: %v", keys(idx)) + } +} + +func TestAST_AsyncFunction(t *testing.T) { + src := ` +async def fetch_data(url: str) -> bytes: + pass +` + result := astExtract(t, "client.py", src, false) + idx := byName(result) + + fn, ok := idx["client.fetch_data"] + if !ok { + t.Fatalf("missing client.fetch_data") + } + if fn.Props["async"] != true { + t.Errorf("fetch_data: async = %v, want true", fn.Props["async"]) + } + if fn.Props["return_type"] != "bytes" { + t.Errorf("fetch_data: return_type = %v, want bytes", fn.Props["return_type"]) + } +} + +func TestAST_DecoratorProps(t *testing.T) { + src := ` +class Repo: + @staticmethod + def from_dict(d): + pass + + @classmethod + def create(cls): + pass + + @property + def name(self): + return self._name +` + result := astExtract(t, "repo.py", src, false) + idx := byName(result) + + cases := []struct { + name string + prop string + want any + }{ + {"repo.Repo.from_dict", "static", true}, + {"repo.Repo.create", "class_method", true}, + {"repo.Repo.name", "property", true}, + } + for _, tc := range cases { + fn, ok := idx[tc.name] + if !ok { + t.Fatalf("missing %q; keys: %v", tc.name, keys(idx)) + } + if fn.Props[tc.prop] != tc.want { + t.Errorf("%s: %s = %v, want %v", tc.name, tc.prop, fn.Props[tc.prop], tc.want) + } + } +} + +func TestAST_SQLAlchemyTable(t *testing.T) { + src := ` +from sqlalchemy import Column, Integer, String +from sqlalchemy.orm import DeclarativeBase + +class Base(DeclarativeBase): + pass + +class Product(Base): + __tablename__ = "products" + id = Column(Integer, primary_key=True) + name = Column(String) +` + result := astExtract(t, "models.py", src, false) + storages := factsByKind(result, facts.KindStorage) + if len(storages) != 1 { + t.Fatalf("expected 1 storage fact, got %d: %v", len(storages), storages) + } + if storages[0].Name != "products" { + t.Errorf("storage name = %q, want products", storages[0].Name) + } + if storages[0].Props["framework"] != "sqlalchemy" { + t.Errorf("storage framework = %v, want sqlalchemy", storages[0].Props["framework"]) + } +} + +func TestAST_ImportEdges(t *testing.T) { + src := ` +import os +from pathlib import Path +from . import utils +` + result := astExtract(t, "mymod.py", src, false) + deps := factsByKind(result, facts.KindDependency) + if len(deps) < 3 { + t.Errorf("expected >= 3 dependency facts, got %d", len(deps)) + } + // Each dep must carry a RelImports relation. + for _, d := range deps { + found := false + for _, r := range d.Relations { + if r.Kind == facts.RelImports { + found = true + } + } + if !found { + t.Errorf("dependency %q missing RelImports relation", d.Name) + } + } +} diff --git a/internal/extractors/pythonextractor/python_test.go b/internal/extractors/pythonextractor/python_test.go index a26ed7a..c32988d 100644 --- a/internal/extractors/pythonextractor/python_test.go +++ b/internal/extractors/pythonextractor/python_test.go @@ -10,22 +10,6 @@ import ( // --- Test helpers --- -// writeAndOpen creates a temp file with the given source content, opens it, -// and returns the open *os.File. The file is closed by the caller. -func writeAndOpen(t *testing.T, filename, src string) *os.File { - t.Helper() - dir := t.TempDir() - path := filepath.Join(dir, filename) - if err := os.WriteFile(path, []byte(src), 0o644); err != nil { - t.Fatal(err) - } - f, err := os.Open(path) - if err != nil { - t.Fatal(err) - } - return f -} - // byName indexes facts by their Name field for easy lookup. func byName(ff []facts.Fact) map[string]facts.Fact { m := make(map[string]facts.Fact, len(ff)) From 96455f423d8d86c61d9d90d7b109f1134d87c688 Mon Sep 17 00:00:00 2001 From: GertL Date: Sun, 14 Jun 2026 18:54:38 +0200 Subject: [PATCH 3/5] Enhance Python AST extractor for Django support and improve route handling - Added handling for assignment nodes in the AST to capture Django urlpatterns at the module level. - Updated the expression statement handler to also process Django urlpatterns. - Enhanced the decorated definition handler to back-fill route facts with the corresponding handler names. - Improved dependency tracking by changing the target format in import statements. - Modified the function handling to distinguish between methods and top-level functions. - Updated tests to use a unified extraction function and removed redundant file handling code. - Cleaned up unused helper tests and ensured consistent handling of class and function definitions. --- internal/extractors/pythonextractor/python.go | 534 +----------------- .../extractors/pythonextractor/python_ast.go | 140 +++-- .../pythonextractor/python_ast_test.go | 3 + .../extractors/pythonextractor/python_test.go | 256 ++------- 4 files changed, 142 insertions(+), 791 deletions(-) diff --git a/internal/extractors/pythonextractor/python.go b/internal/extractors/pythonextractor/python.go index cb5ef24..4f5583a 100644 --- a/internal/extractors/pythonextractor/python.go +++ b/internal/extractors/pythonextractor/python.go @@ -1,7 +1,6 @@ package pythonextractor import ( - "bufio" "context" "io" "io/fs" @@ -122,21 +121,9 @@ func (e *PythonExtractor) Extract(ctx context.Context, repoPath string, files [] return allFacts, nil } -// --- Regex patterns --- +// --- Regex patterns used by the AST walker --- var ( - // classRe matches class declarations. Groups: (indent, name, bases). - classRe = regexp.MustCompile(`^(\s*)class\s+(\w+)\s*(?:\(([^)]*)\))?:`) - - // defRe matches function/method definitions. Groups: (indent, async, name). - defRe = regexp.MustCompile(`^(\s*)(async\s+)?def\s+(\w+)\s*\(`) - - // importRe matches bare import statements. Group: (module). - importRe = regexp.MustCompile(`^\s*import\s+([\w.]+)`) - - // fromImportRe matches from...import statements. Group: (module). - fromImportRe = regexp.MustCompile(`^\s*from\s+([\w.]+)\s+import\s+`) - // routeDecoratorRe matches FastAPI/Starlette route decorators. // Groups: (object, http_method, path). routeDecoratorRe = regexp.MustCompile(`^\s*@([\w.]+)\.(get|post|put|delete|patch|head|options)\s*\(\s*["']([^"']+)["']`) @@ -148,14 +135,6 @@ var ( // Group: (name) e.g. "staticmethod", "app.task". decoratorRe = regexp.MustCompile(`^\s*@([\w.]+)`) - // returnTypeRe extracts return type from a single-line def signature. - // e.g. "def foo(x: int) -> Optional[str]:" - returnTypeRe = regexp.MustCompile(`\)\s*->\s*([\w\[\], |.]+?)\s*:`) - - // returnTypeClosingRe matches the closing paren of a multi-line def with - // a return type annotation. e.g. " ) -> Optional[str]:" - returnTypeClosingRe = regexp.MustCompile(`^\s*\)\s*->\s*([\w\[\], |.]+?)\s*:`) - // apiViewRe matches Django REST Framework @api_view decorators. // Group: (methods_list) — bracket contents, e.g. "'GET', 'POST'" apiViewRe = regexp.MustCompile(`^\s*@(?:[\w.]*\.)?api_view\s*\(\s*\[([^\]]+)\]`) @@ -191,428 +170,6 @@ var ( } ) -// scopeEntry tracks a class nesting level with its indentation. -type scopeEntry struct { - // qualifiedName is the fully-qualified class name (e.g. "dir.Outer.Inner"). - qualifiedName string - // indent is the column indentation of the class keyword. - indent int -} - -// pendingRoute holds a FastAPI route decorator waiting for the handler def. -type pendingRoute struct { - method string - path string - line int -} - -// extractFile parses a single Python file and returns facts. -func extractFile(f *os.File, relFile string, isDjango bool) []facts.Fact { - var result []facts.Fact - dir := filepath.Dir(relFile) - // Python modules are file-based; strip .py to form the module prefix used in - // symbol names (e.g. "app/models/order" for "app/models/order.py"). - module := strings.TrimSuffix(relFile, ".py") - isURLsFile := isDjango && filepath.Base(relFile) == "urls.py" - - scanner := bufio.NewScanner(f) - scanner.Buffer(make([]byte, 0, 256*1024), 1024*1024) - - var ( - lineNum int - scopeStack []scopeEntry - pendingRoutes []pendingRoute - pendingDecorators []string // decorator names since last def/class - pendingApiViewMethods []string // HTTP methods from @api_view - decoratorParenDepth int // open-paren depth inside a multi-line decorator arg list - pendingFuncProps map[string]any // props of last emitted func (for return-type backfill) - pendingFuncLine int - inDocstring bool - docstringQuote string // `"""` or `'''` - ) - - for scanner.Scan() { - lineNum++ - line := scanner.Text() - trimmed := strings.TrimSpace(line) - - // Handle multi-line docstrings / triple-quoted strings. - if inDocstring { - if strings.Contains(line, docstringQuote) { - inDocstring = false - } - continue - } - - // Detect opening of a triple-quoted string. We check after the inDocstring - // block so that a line opening and closing on the same line is handled. - if q, opens := opensTripleQuote(trimmed); opens { - // Count occurrences: if odd number of the quote on this line, we enter - // docstring mode for subsequent lines. - if !closesOnSameLine(trimmed, q) { - inDocstring = true - docstringQuote = q - } - // The line itself is not a declaration, so we can skip to the next line. - // (Triple-quote lines are never class/def/import lines.) - continue - } - - // Skip blank lines and comments. - if trimmed == "" || strings.HasPrefix(trimmed, "#") { - continue - } - - // Determine current line indentation. - indent := lineIndent(line) - - // Pop scope entries that are at the same or deeper indentation level. - scopeStack = popScopes(scopeStack, indent) - - // Expire pending return-type backfill after 20 lines. - if pendingFuncProps != nil && lineNum-pendingFuncLine > 20 { - pendingFuncProps = nil - } - - // Check for return type on the closing line of a multi-line function - // signature, e.g. " ) -> Optional[str]:". The props map is shared with - // the already-emitted fact, so updating it here updates the fact in-place. - if pendingFuncProps != nil { - if rt := returnTypeClosingRe.FindStringSubmatch(line); rt != nil { - pendingFuncProps["return_type"] = strings.TrimSpace(rt[1]) - pendingFuncProps = nil - } - } - - // Class declaration. - if m := classRe.FindStringSubmatch(line); m != nil { - // m[1]=indent, m[2]=name, m[3]=bases (may be empty) - name := m[2] - basesStr := strings.TrimSpace(m[3]) - qualName := buildQualName(module, scopeStack, name) - - props := map[string]any{ - "symbol_kind": facts.SymbolClass, - "language": "python", - } - - rels := []facts.Relation{ - {Kind: facts.RelDeclares, Target: dir}, - } - - bases := splitBases(basesStr) - - // Emit RelImplements for each base class. - for _, base := range bases { - if base != "" { - rels = append(rels, facts.Relation{ - Kind: facts.RelImplements, - Target: base, - }) - } - } - - // Django-specific class classification. - if isDjango { - for _, base := range bases { - bn := lastComponent(base) - if djangoModelBases[bn] { - // Emit KindStorage for the Django-inferred table name. - result = append(result, facts.Fact{ - Kind: facts.KindStorage, - Name: camelToSnake(name), - File: relFile, - Line: lineNum, - Props: map[string]any{ - "storage_kind": "table", - "framework": "django", - "language": "python", - "class": qualName, - }, - Relations: []facts.Relation{ - {Kind: facts.RelDeclares, Target: dir}, - }, - }) - break - } - if djangoCBVBases[bn] { - props["django_component"] = "view" - props["framework"] = "django" - break - } - if djangoSerializerBases[bn] { - props["django_component"] = "serializer" - props["framework"] = "django" - break - } - } - } - - result = append(result, facts.Fact{ - Kind: facts.KindSymbol, - Name: qualName, - File: relFile, - Line: lineNum, - Props: props, - Relations: rels, - }) - - // Push to scope stack so nested members use this class as context. - scopeStack = append(scopeStack, scopeEntry{ - qualifiedName: qualName, - indent: indent, - }) - - // A class declaration closes all pending state. - pendingRoutes = nil - pendingDecorators = nil - pendingApiViewMethods = nil - decoratorParenDepth = 0 - pendingFuncProps = nil - continue - } - - // Function / method definition. - if m := defRe.FindStringSubmatch(line); m != nil { - // m[1]=indent, m[2]=async (may be empty), m[3]=name - isAsync := strings.TrimSpace(m[2]) == "async" - funcName := m[3] - - var fullName string - var symbolKind string - - if len(scopeStack) > 0 { - // We are inside a class — this is a method. - fullName = scopeStack[len(scopeStack)-1].qualifiedName + "." + funcName - symbolKind = facts.SymbolMethod - } else { - // Top-level function. - fullName = module + "." + funcName - symbolKind = facts.SymbolFunc - } - - props := map[string]any{ - "symbol_kind": symbolKind, - "language": "python", - } - if isAsync { - props["async"] = true - } - - // Apply structural props from pending decorators. - for _, dec := range pendingDecorators { - applyDecoratorProps(props, dec) - } - pendingDecorators = nil - decoratorParenDepth = 0 - - // Extract return type from the current line (single-line signature). - if rt := returnTypeRe.FindStringSubmatch(line); rt != nil { - props["return_type"] = strings.TrimSpace(rt[1]) - pendingFuncProps = nil - } else { - // Multi-line signature: backfill return type when closing line appears. - pendingFuncProps = props - pendingFuncLine = lineNum - } - - fact := facts.Fact{ - Kind: facts.KindSymbol, - Name: fullName, - File: relFile, - Line: lineNum, - Props: props, - Relations: []facts.Relation{ - {Kind: facts.RelDeclares, Target: dir}, - }, - } - - // Emit FastAPI route facts. - for _, pr := range pendingRoutes { - result = append(result, facts.Fact{ - Kind: facts.KindRoute, - Name: pr.method + " " + pr.path, - File: relFile, - Line: pr.line, - Props: map[string]any{ - "http_method": pr.method, - "path": pr.path, - "handler": fullName, - "framework": "fastapi", - "language": "python", - }, - }) - } - pendingRoutes = nil - - // Emit Django @api_view route facts. - for _, method := range pendingApiViewMethods { - result = append(result, facts.Fact{ - Kind: facts.KindRoute, - Name: method + " (view) " + fullName, - File: relFile, - Line: lineNum, - Props: map[string]any{ - "http_method": method, - "path": "", - "handler": fullName, - "framework": "django", - "language": "python", - }, - }) - } - pendingApiViewMethods = nil - - result = append(result, fact) - continue - } - - // Route decorator (@router.get("/path"), @app.post("/path"), etc.). - if m := routeDecoratorRe.FindStringSubmatch(line); m != nil { - method := strings.ToUpper(m[2]) - path := m[3] - pendingRoutes = append(pendingRoutes, pendingRoute{ - method: method, - path: path, - line: lineNum, - }) - continue - } - - // Any decorator line (@...). - if strings.HasPrefix(trimmed, "@") { - // Django @api_view(['GET', 'POST']) — parse HTTP methods. - if isDjango { - if m := apiViewRe.FindStringSubmatch(line); m != nil { - methods := httpMethodWordRe.FindAllString(m[1], -1) - pendingApiViewMethods = append(pendingApiViewMethods, methods...) - decoratorParenDepth = 0 // api_view always single-line - continue - } - } - - // Capture decorator name for structural symbol props. - if m := decoratorRe.FindStringSubmatch(line); m != nil { - pendingDecorators = append(pendingDecorators, m[1]) - } - // Track open parens so multi-line decorator args don't clear pending state. - decoratorParenDepth = strings.Count(line, "(") - strings.Count(line, ")") - if decoratorParenDepth < 0 { - decoratorParenDepth = 0 - } - continue - } - - // Inside a multi-line decorator argument list — update depth and wait. - if decoratorParenDepth > 0 { - decoratorParenDepth += strings.Count(line, "(") - strings.Count(line, ")") - if decoratorParenDepth < 0 { - decoratorParenDepth = 0 - } - continue - } - - // Any non-decorator, non-def line clears all pending state. - pendingRoutes = nil - pendingDecorators = nil - pendingApiViewMethods = nil - - // Django URL patterns in urls.py files. - if isURLsFile { - if m := urlPathRe.FindStringSubmatch(line); m != nil { - result = append(result, facts.Fact{ - Kind: facts.KindRoute, - Name: "* " + m[1], - File: relFile, - Line: lineNum, - Props: map[string]any{ - "path": m[1], - "handler": m[2], - "framework": "django", - "language": "python", - }, - }) - } - } - - // Import: `import foo.bar` - if m := importRe.FindStringSubmatch(line); m != nil { - importPath := m[1] - result = append(result, facts.Fact{ - Kind: facts.KindDependency, - Name: module + " -> " + importPath, - File: relFile, - Line: lineNum, - Props: map[string]any{ - "language": "python", - }, - Relations: []facts.Relation{ - {Kind: facts.RelImports, Target: importPath}, - }, - }) - continue - } - - // Import: `from foo.bar import ...` - if m := fromImportRe.FindStringSubmatch(line); m != nil { - importPath := m[1] - result = append(result, facts.Fact{ - Kind: facts.KindDependency, - Name: module + " -> " + importPath, - File: relFile, - Line: lineNum, - Props: map[string]any{ - "language": "python", - "from": true, - }, - Relations: []facts.Relation{ - {Kind: facts.RelImports, Target: importPath}, - }, - }) - continue - } - - // SQLAlchemy table name: `__tablename__ = "tbl"` - if m := tableNameRe.FindStringSubmatch(line); m != nil { - tableName := m[1] - - // Determine the owning class from the scope stack. - ownerClass := "" - if len(scopeStack) > 0 { - ownerClass = scopeStack[len(scopeStack)-1].qualifiedName - } - - props := map[string]any{ - "storage_kind": "table", - "framework": "sqlalchemy", - "language": "python", - } - if ownerClass != "" { - props["class"] = ownerClass - } - - result = append(result, facts.Fact{ - Kind: facts.KindStorage, - Name: tableName, - File: relFile, - Line: lineNum, - Props: props, - Relations: []facts.Relation{ - {Kind: facts.RelDeclares, Target: dir}, - }, - }) - continue - } - } - - if err := scanner.Err(); err != nil { - log.Printf("[python-extractor] scanner error in %s: %v", relFile, err) - } - - return result -} - -// --- Helpers --- // applyDecoratorProps sets structural boolean props on a symbol based on a // decorator name. Only well-known structural decorators produce props; unknown @@ -684,95 +241,6 @@ func lastComponent(name string) string { return name } -// buildQualName constructs a qualified name like "module.Outer.Inner.Name". -// module is the file-based module path (e.g. "app/models/order" for "app/models/order.py"). -func buildQualName(module string, stack []scopeEntry, name string) string { - if len(stack) == 0 { - return module + "." + name - } - return stack[len(stack)-1].qualifiedName + "." + name -} - -// popScopes removes scope entries at or deeper than the given indentation. -func popScopes(stack []scopeEntry, indent int) []scopeEntry { - for len(stack) > 0 && stack[len(stack)-1].indent >= indent { - stack = stack[:len(stack)-1] - } - return stack -} - -// lineIndent returns the number of leading spaces in a line. -func lineIndent(line string) int { - count := 0 - for _, ch := range line { - if ch == ' ' { - count++ - } else if ch == '\t' { - count += 4 // treat tab as 4 spaces - } else { - break - } - } - return count -} - -// opensTripleQuote checks if a trimmed line starts a triple-quoted string and -// returns the quote style (`"""` or `'''`) and whether it opens one. -func opensTripleQuote(trimmed string) (string, bool) { - for _, q := range []string{`"""`, `'''`} { - if strings.Contains(trimmed, q) { - return q, true - } - } - return "", false -} - -// closesOnSameLine returns true if the triple quote appears an even number of -// times on the line (opened and closed on the same line). -func closesOnSameLine(trimmed, q string) bool { - count := strings.Count(trimmed, q) - return count >= 2 -} - -// splitBases splits a Python base class list by comma, respecting bracket nesting -// so that generic types like `Generic[T]` or `CRUDBase[Model, Schema]` are kept -// as a single token. -func splitBases(s string) []string { - var result []string - depth := 0 - start := 0 - for i, ch := range s { - switch ch { - case '[', '(': - depth++ - case ']', ')': - depth-- - case ',': - if depth == 0 { - if t := strings.TrimSpace(s[start:i]); t != "" { - result = append(result, stripGeneric(t)) - } - start = i + 1 - } - } - } - if t := strings.TrimSpace(s[start:]); t != "" { - result = append(result, stripGeneric(t)) - } - return result -} - -// stripGeneric removes generic type parameters from a base class name. -// e.g. "Generic[T]" → "Generic", "CRUDBase[Model, Schema]" → "CRUDBase". -func stripGeneric(s string) string { - if before, _, ok := strings.Cut(s, "["); ok { - return strings.TrimSpace(before) - } - if before, _, ok := strings.Cut(s, "("); ok { - return strings.TrimSpace(before) - } - return strings.TrimSpace(s) -} // isPythonFile returns true if the file has a .py extension. func isPythonFile(path string) bool { diff --git a/internal/extractors/pythonextractor/python_ast.go b/internal/extractors/pythonextractor/python_ast.go index dede3fc..6b27637 100644 --- a/internal/extractors/pythonextractor/python_ast.go +++ b/internal/extractors/pythonextractor/python_ast.go @@ -121,7 +121,11 @@ func (w *pyWalker) walkStatement(node *sitter.Node) { w.handleDecoratedDefinition(node) case "expression_statement": // __tablename__ = "foo" (SQLAlchemy) lives here at class body level. + // urlpatterns = [...] (Django) lives at module level. w.handleExprStatement(node) + case "assignment": + // tree-sitter may parse assignments as "assignment" nodes at module level. + w.handleAssignment(node) case "block": for i := uint(0); i < uint(node.ChildCount()); i++ { w.walkStatement(node.Child(i)) @@ -148,7 +152,7 @@ func (w *pyWalker) handleImport(node *sitter.Node) { } else { name = pyText(c, w.src) } - target := w.dir + " -> " + name + target := w.module + " -> " + name w.out = append(w.out, facts.Fact{ Kind: facts.KindDependency, Name: target, @@ -184,13 +188,13 @@ func (w *pyWalker) handleFromImport(node *sitter.Node) { isRelative := strings.HasPrefix(moduleName, ".") || strings.HasPrefix(pyText(node, w.src), "from .") - target := w.dir + " -> " + moduleName + target := w.module + " -> " + moduleName w.out = append(w.out, facts.Fact{ Kind: facts.KindDependency, Name: target, File: w.relFile, Line: int(node.StartPosition().Row) + 1, - Props: map[string]any{"language": "python"}, + Props: map[string]any{"language": "python", "from": true}, Relations: []facts.Relation{ {Kind: facts.RelImports, Target: moduleName}, }, @@ -248,6 +252,8 @@ func (w *pyWalker) setImport(local, target string) { func (w *pyWalker) handleDecoratedDefinition(node *sitter.Node) { var decorators []string var pendingApiViewMethods []string + // pendingRoutes holds route facts emitted from decorators before we see the handler name. + var pendingRouteFacts []*facts.Fact for i := uint(0); i < uint(node.ChildCount()); i++ { c := node.Child(i) @@ -256,17 +262,21 @@ func (w *pyWalker) handleDecoratedDefinition(node *sitter.Node) { text := pyText(c, w.src) // FastAPI / Starlette route decorator. if m := routeDecoratorRe.FindStringSubmatch(text); m != nil { - w.out = append(w.out, facts.Fact{ + method := strings.ToUpper(m[2]) + path := m[3] + rf := facts.Fact{ Kind: facts.KindRoute, - Name: strings.ToUpper(m[2]) + " " + m[3], + Name: method + " " + path, File: w.relFile, Line: int(c.StartPosition().Row) + 1, Props: map[string]any{ - "method": strings.ToUpper(m[2]), - "path": m[3], - "framework": "fastapi", + "http_method": method, + "path": path, + "framework": "fastapi", }, - }) + } + w.out = append(w.out, rf) + pendingRouteFacts = append(pendingRouteFacts, &w.out[len(w.out)-1]) continue } // DRF @api_view(['GET','POST']). @@ -283,9 +293,13 @@ func (w *pyWalker) handleDecoratedDefinition(node *sitter.Node) { case "function_definition": w.handleFunction(c, decorators) + handlerName := w.module + "." + w.qualify(pyFuncName(c, w.src)) + // Back-fill handler into pending FastAPI route facts. + for _, rf := range pendingRouteFacts { + rf.Props["handler"] = handlerName + } // @api_view routes — emit after we know the handler name. if len(pendingApiViewMethods) > 0 { - handlerName := w.module + "." + w.qualify(pyFuncName(c, w.src)) for _, meth := range pendingApiViewMethods { w.out = append(w.out, facts.Fact{ Kind: facts.KindRoute, @@ -293,9 +307,9 @@ func (w *pyWalker) handleDecoratedDefinition(node *sitter.Node) { File: w.relFile, Line: int(c.StartPosition().Row) + 1, Props: map[string]any{ - "method": meth, - "framework": "django", - "handler": handlerName, + "http_method": meth, + "framework": "django", + "handler": handlerName, }, }) } @@ -327,10 +341,23 @@ func (w *pyWalker) handleClass(node *sitter.Node, decorators []string) { if args := node.ChildByFieldName("superclasses"); args != nil { for i := uint(0); i < uint(args.ChildCount()); i++ { c := args.Child(i) - if c.Kind() == "identifier" || c.Kind() == "attribute" { + switch c.Kind() { + case "identifier": + base := pyText(c, w.src) + bases = append(bases, base) + rels = append(rels, facts.Relation{Kind: facts.RelImplements, Target: base}) + case "attribute": base := pyText(c, w.src) bases = append(bases, base) rels = append(rels, facts.Relation{Kind: facts.RelImplements, Target: base}) + case "subscript": + // Generic base: CRUDBase[ModelType, IdType] — strip the type params. + valueNode := c.ChildByFieldName("value") + if valueNode != nil { + base := pyText(valueNode, w.src) + bases = append(bases, base) + rels = append(rels, facts.Relation{Kind: facts.RelImplements, Target: base}) + } } } } @@ -372,27 +399,6 @@ func (w *pyWalker) handleClass(node *sitter.Node, decorators []string) { } } - // Django urls.py: emit route facts from path()/re_path() calls in the class body. - if w.isDjango && filepath.Base(w.relFile) == "urls.py" { - bodyNode := node.ChildByFieldName("body") - if bodyNode != nil { - bodyText := pyText(bodyNode, w.src) - for _, m := range urlPathRe.FindAllStringSubmatch(bodyText, -1) { - w.out = append(w.out, facts.Fact{ - Kind: facts.KindRoute, - Name: "* " + m[1], - File: w.relFile, - Line: int(node.StartPosition().Row) + 1, - Props: map[string]any{ - "path": m[1], - "handler": m[2], - "framework": "django", - }, - }) - } - } - } - f := facts.Fact{ Kind: facts.KindSymbol, Name: qualName, @@ -424,8 +430,14 @@ func (w *pyWalker) handleFunction(node *sitter.Node, decorators []string) { name := pyText(nameNode, w.src) qualName := w.module + "." + w.qualify(name) + // Determine if this is a method (inside a class) or a top-level function. + symbolKind := facts.SymbolFunc + if len(w.typeStack) > 0 { + symbolKind = facts.SymbolMethod + } + props := map[string]any{ - "symbol_kind": facts.SymbolFunc, + "symbol_kind": symbolKind, "language": "python", } if len(w.typeStack) > 0 { @@ -473,11 +485,17 @@ func (w *pyWalker) handleFunction(node *sitter.Node, decorators []string) { w.popOwner() } -// handleExprStatement checks for SQLAlchemy __tablename__ assignments. +// handleExprStatement checks for SQLAlchemy __tablename__ assignments and +// Django urlpatterns at module/class level. func (w *pyWalker) handleExprStatement(node *sitter.Node) { text := pyText(node, w.src) if m := tableNameRe.FindStringSubmatch(text); m != nil { - w.out = append(w.out, facts.Fact{ + // Find the enclosing class name for the storage fact. + className := "" + if len(w.typeStack) > 0 { + className = w.module + "." + w.enclosingType() + } + sf := facts.Fact{ Kind: facts.KindStorage, Name: m[1], File: w.relFile, @@ -486,7 +504,51 @@ func (w *pyWalker) handleExprStatement(node *sitter.Node) { "storage_kind": "table", "framework": "sqlalchemy", }, - }) + } + if className != "" { + sf.Props["class"] = className + sf.Relations = []facts.Relation{{Kind: facts.RelDeclares, Target: w.dir}} + } + w.out = append(w.out, sf) + return + } + // Django urls.py: urlpatterns = [...]. + if w.isDjango && filepath.Base(w.relFile) == "urls.py" { + for _, m := range urlPathRe.FindAllStringSubmatch(text, -1) { + w.out = append(w.out, facts.Fact{ + Kind: facts.KindRoute, + Name: "* " + m[1], + File: w.relFile, + Line: int(node.StartPosition().Row) + 1, + Props: map[string]any{ + "path": m[1], + "handler": m[2], + "framework": "django", + }, + }) + } + } +} + +// handleAssignment handles module-level assignment statements (tree-sitter +// sometimes emits these as "assignment" nodes rather than "expression_statement"). +func (w *pyWalker) handleAssignment(node *sitter.Node) { + text := pyText(node, w.src) + // Django urls.py: urlpatterns = [...]. + if w.isDjango && filepath.Base(w.relFile) == "urls.py" { + for _, m := range urlPathRe.FindAllStringSubmatch(text, -1) { + w.out = append(w.out, facts.Fact{ + Kind: facts.KindRoute, + Name: "* " + m[1], + File: w.relFile, + Line: int(node.StartPosition().Row) + 1, + Props: map[string]any{ + "path": m[1], + "handler": m[2], + "framework": "django", + }, + }) + } } } diff --git a/internal/extractors/pythonextractor/python_ast_test.go b/internal/extractors/pythonextractor/python_ast_test.go index e68e9e0..689c80e 100644 --- a/internal/extractors/pythonextractor/python_ast_test.go +++ b/internal/extractors/pythonextractor/python_ast_test.go @@ -13,6 +13,9 @@ func astExtract(t *testing.T, filename, src string, isDjango bool) []facts.Fact t.Helper() dir := t.TempDir() path := filepath.Join(dir, filename) + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatal(err) + } if err := os.WriteFile(path, []byte(src), 0o644); err != nil { t.Fatal(err) } diff --git a/internal/extractors/pythonextractor/python_test.go b/internal/extractors/pythonextractor/python_test.go index c32988d..156f536 100644 --- a/internal/extractors/pythonextractor/python_test.go +++ b/internal/extractors/pythonextractor/python_test.go @@ -61,11 +61,8 @@ class Order: def calculate(self): return self.total * 1.2 ` - f := writeAndOpen(t, "order.py", src) - defer f.Close() - relFile := "app/models/order.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) // Class fact: module.ClassName @@ -107,11 +104,8 @@ def helper(x, y): async def fetch_data(url): pass ` - f := writeAndOpen(t, "utils.py", src) - defer f.Close() - relFile := "services/utils.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) helperName := mod(relFile) + ".helper" @@ -142,11 +136,8 @@ class VespaSink(EmbeddingsSink): def send(self, data): pass ` - f := writeAndOpen(t, "vespa_sink.py", src) - defer f.Close() - relFile := "sinks/vespa_sink.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) clsName := mod(relFile) + ".VespaSink" @@ -167,11 +158,8 @@ class FeatureGroup(Base, TimestampMixin): def validate(self): pass ` - f := writeAndOpen(t, "feature_group.py", src) - defer f.Close() - relFile := "db/models/feature_group.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) clsName := mod(relFile) + ".FeatureGroup" @@ -192,11 +180,8 @@ func TestExtractFile_ClassInheritance_GenericBase(t *testing.T) { class CRUDEntity(CRUDBase[ModelType, IdType]): pass ` - f := writeAndOpen(t, "crud.py", src) - defer f.Close() - relFile := "db/crud/crud.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) clsName := mod(relFile) + ".CRUDEntity" @@ -220,11 +205,8 @@ class Outer: def outer_method(self): pass ` - f := writeAndOpen(t, "nested.py", src) - defer f.Close() - relFile := "pkg/nested.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) outerName := mod(relFile) + ".Outer" @@ -265,11 +247,8 @@ class Recommender: async def recommend(self, user_id): pass ` - f := writeAndOpen(t, "recommender.py", src) - defer f.Close() - relFile := "services/recommender.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) methodName := mod(relFile) + ".Recommender.recommend" @@ -291,11 +270,8 @@ import logging import os import fastapi ` - f := writeAndOpen(t, "app.py", src) - defer f.Close() - relFile := "myapp/app.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) for _, target := range []string{"logging", "os", "fastapi"} { @@ -320,11 +296,8 @@ from fastapi import APIRouter, Depends from query_recommender.models.filters import VespaSearchFilters from .base import Base ` - f := writeAndOpen(t, "routes.py", src) - defer f.Close() - relFile := "routes/routes.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) cases := []struct { @@ -361,11 +334,8 @@ router = APIRouter() async def health_check(): return {"status": "ok"} ` - f := writeAndOpen(t, "health.py", src) - defer f.Close() - relFile := "routes/health.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) routes := factsByKind(result, facts.KindRoute) if len(routes) != 1 { @@ -398,11 +368,8 @@ router = APIRouter() async def post_recommend_v2(body: RecommendV2Body) -> RecommendV2Response: pass ` - f := writeAndOpen(t, "recommend.py", src) - defer f.Close() - relFile := "routes/recommend.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) routes := factsByKind(result, facts.KindRoute) if len(routes) != 1 { @@ -433,10 +400,7 @@ async def create_item(): async def delete_item(id: int): pass ` - f := writeAndOpen(t, "items.py", src) - defer f.Close() - - result := extractFile(f, "routes/items.py", false) + result := astExtract(t, "routes/items.py", src, false) routes := factsByKind(result, facts.KindRoute) if len(routes) != 3 { @@ -462,10 +426,7 @@ router = APIRouter() async def login(): pass ` - f := writeAndOpen(t, "auth.py", src) - defer f.Close() - - result := extractFile(f, "routes/auth.py", false) + result := astExtract(t, "routes/auth.py", src, false) routes := factsByKind(result, facts.KindRoute) if len(routes) != 1 { @@ -486,11 +447,8 @@ class FeatureGroup(Base): id: Mapped[int] name: Mapped[str] ` - f := writeAndOpen(t, "feature_group.py", src) - defer f.Close() - relFile := "db/models/feature_group.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) storages := factsByKind(result, facts.KindStorage) if len(storages) != 1 { @@ -529,11 +487,8 @@ class MyService: def real_method(self): pass ` - f := writeAndOpen(t, "service.py", src) - defer f.Close() - relFile := "services/service.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) // fake_def and FakeClass inside the docstring must NOT appear. @@ -563,11 +518,8 @@ class Validator: def validate(self, value): pass ` - f := writeAndOpen(t, "validator.py", src) - defer f.Close() - relFile := "pkg/validator.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) validateName := mod(relFile) + ".Validator.validate" @@ -581,11 +533,8 @@ func TestExtractFile_LineNumbers(t *testing.T) { def bar(self): pass ` - f := writeAndOpen(t, "foo.py", src) - defer f.Close() - relFile := "pkg/foo.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) clsName := mod(relFile) + ".Foo" @@ -614,11 +563,8 @@ func TestExtractFile_ClassWithoutBases(t *testing.T) { class Foo: pass ` - f := writeAndOpen(t, "foo.py", src) - defer f.Close() - relFile := "pkg/foo.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) clsName := mod(relFile) + ".Foo" @@ -658,11 +604,8 @@ async def post_recommend_v2( ) -> None: pass ` - f := writeAndOpen(t, "recommend_v2.py", src) - defer f.Close() - relFile := "query_recommender/routes/recommend_v2.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) routes := factsByKind(result, facts.KindRoute) if len(routes) != 1 { @@ -711,11 +654,8 @@ class Entity(Base): def __repr__(self) -> str: return f"" ` - f := writeAndOpen(t, "entity.py", src) - defer f.Close() - relFile := "db/models/entity.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) // Class with Base inheritance. @@ -821,77 +761,6 @@ func TestDetect_GoRepo(t *testing.T) { } } -// --- Helper unit tests --- - -func TestSplitBases_Simple(t *testing.T) { - cases := []struct { - input string - want []string - }{ - {"Base", []string{"Base"}}, - {"Base, Mixin", []string{"Base", "Mixin"}}, - {"CRUDBase[Model, Schema]", []string{"CRUDBase"}}, - {"Generic[T], Protocol", []string{"Generic", "Protocol"}}, - } - for _, tc := range cases { - got := splitBases(tc.input) - if len(got) != len(tc.want) { - t.Errorf("splitBases(%q): got %v (len %d), want %v (len %d)", - tc.input, got, len(got), tc.want, len(tc.want)) - continue - } - for i, g := range got { - if g != tc.want[i] { - t.Errorf("splitBases(%q)[%d]: got %q, want %q", tc.input, i, g, tc.want[i]) - } - } - } -} - -func TestSplitBases_Empty(t *testing.T) { - got := splitBases("") - if len(got) != 0 { - t.Errorf("splitBases(%q): got %v, want empty", "", got) - } -} - -func TestPopScopes(t *testing.T) { - stack := []scopeEntry{ - {qualifiedName: "pkg.Outer", indent: 0}, - {qualifiedName: "pkg.Outer.Inner", indent: 4}, - } - - // A line at indent=4 should pop Inner (4 >= 4) but keep Outer (0 < 4). - got := popScopes(stack, 4) - if len(got) != 1 || got[0].qualifiedName != "pkg.Outer" { - t.Errorf("popScopes at indent=4: got %v, want [pkg.Outer]", got) - } - - // A line at indent=0 should pop everything. - got = popScopes(stack, 0) - if len(got) != 0 { - t.Errorf("popScopes at indent=0: got %v, want []", got) - } -} - -func TestLineIndent(t *testing.T) { - cases := []struct { - line string - want int - }{ - {"class Foo:", 0}, - {" def bar(self):", 4}, - {" pass", 8}, - {"\t\tpass", 8}, // tab = 4 spaces - {"", 0}, - } - for _, tc := range cases { - got := lineIndent(tc.line) - if got != tc.want { - t.Errorf("lineIndent(%q) = %d, want %d", tc.line, got, tc.want) - } - } -} // keys returns map keys sorted for deterministic error messages. func keys(m map[string]facts.Fact) []string { @@ -915,11 +784,8 @@ class Config: def items(self): return [] ` - f := writeAndOpen(t, "config.py", src) - defer f.Close() - relFile := "app/config.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) for _, methodName := range []string{ @@ -943,11 +809,8 @@ class Utils: def parse(value): return int(value) ` - f := writeAndOpen(t, "utils.py", src) - defer f.Close() - relFile := "pkg/utils.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) methName := mod(relFile) + ".Utils.parse" @@ -970,11 +833,8 @@ class Repo: def from_env(cls): pass ` - f := writeAndOpen(t, "repo.py", src) - defer f.Close() - relFile := "db/repo.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) methName := mod(relFile) + ".Repo.from_env" @@ -1002,11 +862,8 @@ class Base(ABC): def concrete(self): pass ` - f := writeAndOpen(t, "base.py", src) - defer f.Close() - relFile := "core/base.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) executeName := mod(relFile) + ".Base.execute" @@ -1040,11 +897,8 @@ class Base: def create(cls): pass ` - f := writeAndOpen(t, "base.py", src) - defer f.Close() - relFile := "core/base.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) methName := mod(relFile) + ".Base.create" @@ -1067,11 +921,8 @@ func TestExtractFile_Task_Bare(t *testing.T) { def process_records(): pass ` - f := writeAndOpen(t, "tasks.py", src) - defer f.Close() - relFile := "jobs/tasks.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) fnName := mod(relFile) + ".process_records" @@ -1096,11 +947,8 @@ from celery import shared_task def send_welcome_email(user_id: int) -> None: pass ` - f := writeAndOpen(t, "email_tasks.py", src) - defer f.Close() - relFile := "notifications/email_tasks.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) fnName := mod(relFile) + ".send_welcome_email" @@ -1127,11 +975,8 @@ func TestExtractFile_MultiLineDecorator(t *testing.T) { def retry_task(self): pass ` - f := writeAndOpen(t, "tasks.py", src) - defer f.Close() - relFile := "jobs/tasks.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) fnName := mod(relFile) + ".retry_task" @@ -1157,11 +1002,8 @@ def get_count() -> int: def no_annotation(): pass ` - f := writeAndOpen(t, "funcs.py", src) - defer f.Close() - relFile := "pkg/funcs.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) cases := []struct { @@ -1199,11 +1041,8 @@ def find_user() -> Optional[str]: def get_items() -> list[str] | None: pass ` - f := writeAndOpen(t, "service.py", src) - defer f.Close() - relFile := "svc/service.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) cases := []struct{ name, want string }{ @@ -1231,11 +1070,8 @@ def create_handler( ) -> Optional[str]: pass ` - f := writeAndOpen(t, "handler.py", src) - defer f.Close() - relFile := "api/handler.py" - result := extractFile(f, relFile, false) + result := astExtract(t, relFile, src, false) idx := byName(result) fnName := mod(relFile) + ".create_handler" @@ -1260,11 +1096,8 @@ class Order(models.Model): class UserProfile(models.Model): user = models.OneToOneField('User', on_delete=models.CASCADE) ` - f := writeAndOpen(t, "models.py", src) - defer f.Close() - relFile := "shop/models.py" - result := extractFile(f, relFile, true) + result := astExtract(t, relFile, src, true) storages := factsByKind(result, facts.KindStorage) if len(storages) != 2 { @@ -1303,11 +1136,8 @@ class OrderView(APIView): def get(self, request): pass ` - f := writeAndOpen(t, "views.py", src) - defer f.Close() - relFile := "shop/views.py" - result := extractFile(f, relFile, true) + result := astExtract(t, relFile, src, true) idx := byName(result) clsName := mod(relFile) + ".OrderView" @@ -1332,11 +1162,8 @@ class OrderSerializer(serializers.ModelSerializer): model = Order fields = '__all__' ` - f := writeAndOpen(t, "serializers.py", src) - defer f.Close() - relFile := "shop/serializers.py" - result := extractFile(f, relFile, true) + result := astExtract(t, relFile, src, true) idx := byName(result) clsName := mod(relFile) + ".OrderSerializer" @@ -1364,10 +1191,7 @@ urlpatterns = [ ] ` // File must be named urls.py for Django URL extraction. - f := writeAndOpen(t, "urls.py", src) - defer f.Close() - - result := extractFile(f, "shop/urls.py", true) + result := astExtract(t, "shop/urls.py", src, true) routes := factsByKind(result, facts.KindRoute) if len(routes) != 3 { @@ -1394,10 +1218,7 @@ urlpatterns = [ path('orders/', views.OrderListView.as_view()), ] ` - f := writeAndOpen(t, "routing.py", src) - defer f.Close() - - result := extractFile(f, "shop/routing.py", true) + result := astExtract(t, "shop/routing.py", src, true) routes := factsByKind(result, facts.KindRoute) if len(routes) != 0 { t.Errorf("expected no routes in non-urls.py file, got %d", len(routes)) @@ -1416,11 +1237,8 @@ def order_list(request): def order_detail(request, pk): pass ` - f := writeAndOpen(t, "views.py", src) - defer f.Close() - relFile := "shop/views.py" - result := extractFile(f, relFile, true) + result := astExtract(t, relFile, src, true) routes := factsByKind(result, facts.KindRoute) // order_list has GET+POST → 2 routes; order_detail has GET → 1 route. From 9aafef3abff442cabf3e475bc5730650f18ba348 Mon Sep 17 00:00:00 2001 From: GertL Date: Mon, 15 Jun 2026 07:24:18 +0200 Subject: [PATCH 4/5] Fix stale pointer bug in pyWalker and skip @overload stubs Replace ownerStack/pendingRouteIndices pointer captures with slice indices to prevent silently lost call edges and handler back-fills when w.out reallocates its backing array during body traversal. Skip @overload-decorated function definitions to avoid emitting duplicate symbol facts for type-checker-only stubs that have no runtime behaviour. Co-Authored-By: Claude Sonnet 4.6 --- .../extractors/pythonextractor/python_ast.go | 52 +++++++++++++------ 1 file changed, 35 insertions(+), 17 deletions(-) diff --git a/internal/extractors/pythonextractor/python_ast.go b/internal/extractors/pythonextractor/python_ast.go index 6b27637..c7fc05f 100644 --- a/internal/extractors/pythonextractor/python_ast.go +++ b/internal/extractors/pythonextractor/python_ast.go @@ -50,9 +50,11 @@ type pyWalker struct { // typeStack holds enclosing class names so methods get qualified names. typeStack []string - // ownerStack: top element is the fact that receives RelCalls / RelInstantiates - // discovered while walking its body. - ownerStack []*facts.Fact + // ownerStack: top element is the index into w.out of the fact that receives + // RelCalls / RelInstantiates discovered while walking its body. Indices are + // used instead of pointers because appending to w.out can reallocate the + // backing array, invalidating any previously captured pointer. + ownerStack []int // importMap maps a local name to its canonical fact target (empty = external). importMap map[string]string @@ -62,13 +64,13 @@ type pyWalker struct { methodSets []map[string]bool } -func (w *pyWalker) pushOwner(f *facts.Fact) { w.ownerStack = append(w.ownerStack, f) } +func (w *pyWalker) pushOwner(idx int) { w.ownerStack = append(w.ownerStack, idx) } func (w *pyWalker) popOwner() { w.ownerStack = w.ownerStack[:len(w.ownerStack)-1] } func (w *pyWalker) currentOwner() *facts.Fact { if len(w.ownerStack) == 0 { return nil } - return w.ownerStack[len(w.ownerStack)-1] + return &w.out[w.ownerStack[len(w.ownerStack)-1]] } func (w *pyWalker) enclosingType() string { return strings.Join(w.typeStack, ".") } @@ -252,8 +254,10 @@ func (w *pyWalker) setImport(local, target string) { func (w *pyWalker) handleDecoratedDefinition(node *sitter.Node) { var decorators []string var pendingApiViewMethods []string - // pendingRoutes holds route facts emitted from decorators before we see the handler name. - var pendingRouteFacts []*facts.Fact + // pendingRouteIndices holds w.out indices of route facts emitted from + // decorators before we see the handler name. Indices are used (not pointers) + // because subsequent appends to w.out may reallocate its backing array. + var pendingRouteIndices []int for i := uint(0); i < uint(node.ChildCount()); i++ { c := node.Child(i) @@ -264,7 +268,7 @@ func (w *pyWalker) handleDecoratedDefinition(node *sitter.Node) { if m := routeDecoratorRe.FindStringSubmatch(text); m != nil { method := strings.ToUpper(m[2]) path := m[3] - rf := facts.Fact{ + w.out = append(w.out, facts.Fact{ Kind: facts.KindRoute, Name: method + " " + path, File: w.relFile, @@ -274,9 +278,8 @@ func (w *pyWalker) handleDecoratedDefinition(node *sitter.Node) { "path": path, "framework": "fastapi", }, - } - w.out = append(w.out, rf) - pendingRouteFacts = append(pendingRouteFacts, &w.out[len(w.out)-1]) + }) + pendingRouteIndices = append(pendingRouteIndices, len(w.out)-1) continue } // DRF @api_view(['GET','POST']). @@ -292,11 +295,16 @@ func (w *pyWalker) handleDecoratedDefinition(node *sitter.Node) { } case "function_definition": + // @overload stubs are type-checker-only annotations with no runtime + // body — skip them to avoid duplicate symbol facts. + if hasDecorator(decorators, "overload") { + continue + } w.handleFunction(c, decorators) handlerName := w.module + "." + w.qualify(pyFuncName(c, w.src)) // Back-fill handler into pending FastAPI route facts. - for _, rf := range pendingRouteFacts { - rf.Props["handler"] = handlerName + for _, idx := range pendingRouteIndices { + w.out[idx].Props["handler"] = handlerName } // @api_view routes — emit after we know the handler name. if len(pendingApiViewMethods) > 0 { @@ -409,8 +417,7 @@ func (w *pyWalker) handleClass(node *sitter.Node, decorators []string) { } w.out = append(w.out, f) - owner := &w.out[len(w.out)-1] - w.pushOwner(owner) + w.pushOwner(len(w.out) - 1) bodyNode := node.ChildByFieldName("body") w.pushType(name, collectPyMethodNames(bodyNode, w.src)) @@ -477,8 +484,7 @@ func (w *pyWalker) handleFunction(node *sitter.Node, decorators []string) { } w.out = append(w.out, f) - owner := &w.out[len(w.out)-1] - w.pushOwner(owner) + w.pushOwner(len(w.out) - 1) if bodyNode := node.ChildByFieldName("body"); bodyNode != nil { w.walkForCalls(bodyNode) } @@ -671,6 +677,18 @@ func collectPyMethodNames(body *sitter.Node, src []byte) map[string]bool { return methods } +// hasDecorator reports whether any name in decorators has last as its +// last dot-separated component (e.g. "overload" matches both "overload" +// and "typing.overload"). +func hasDecorator(decorators []string, last string) bool { + for _, d := range decorators { + if lastComponent(d) == last { + return true + } + } + return false +} + func pyFuncName(node *sitter.Node, src []byte) string { if n := node.ChildByFieldName("name"); n != nil { return pyText(n, src) From 4654614c2977eeb9397da5b0c7b700ba0cc0e3c1 Mon Sep 17 00:00:00 2001 From: GertL Date: Tue, 16 Jun 2026 06:09:31 +0200 Subject: [PATCH 5/5] Update tree-sitter-python dependency to v0.23.6 in go.mod and go.sum --- go.mod | 2 +- go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index d8266ba..21e5b42 100644 --- a/go.mod +++ b/go.mod @@ -6,7 +6,7 @@ require ( github.com/modelcontextprotocol/go-sdk v1.3.0 github.com/tree-sitter-grammars/tree-sitter-kotlin v1.1.0 github.com/tree-sitter/go-tree-sitter v0.24.0 - github.com/tree-sitter/tree-sitter-python v0.21.1-0.20240818005537-55a9b8a4fbfb + github.com/tree-sitter/tree-sitter-python v0.23.6 github.com/tree-sitter/tree-sitter-typescript v0.23.2 gopkg.in/yaml.v3 v3.0.1 ) diff --git a/go.sum b/go.sum index 49e74aa..8d4b5b1 100644 --- a/go.sum +++ b/go.sum @@ -36,8 +36,8 @@ github.com/tree-sitter/tree-sitter-json v0.21.1-0.20240818005659-bdd69eb8c8a5 h1 github.com/tree-sitter/tree-sitter-json v0.21.1-0.20240818005659-bdd69eb8c8a5/go.mod h1:GbMKRjLfk0H+PI7nLi1Sx5lHf5wCpLz9al8tQYSxpEk= github.com/tree-sitter/tree-sitter-php v0.22.9-0.20240819002312-a552625b56c1 h1:ZXZMDwE+IhUtGug4Brv6NjJWUU3rfkZBKpemf6RY8/g= github.com/tree-sitter/tree-sitter-php v0.22.9-0.20240819002312-a552625b56c1/go.mod h1:UKCLuYnJ312Mei+3cyTmGOHzn0YAnaPRECgJmHtzrqs= -github.com/tree-sitter/tree-sitter-python v0.21.1-0.20240818005537-55a9b8a4fbfb h1:EXEM82lFM7JjJb6qiKZXkpIDaCcbV2obNn82ghwj9lw= -github.com/tree-sitter/tree-sitter-python v0.21.1-0.20240818005537-55a9b8a4fbfb/go.mod h1:lXCF1nGG5Dr4J3BTS0ObN4xJCCICiSu/b+Xe/VqMV7g= +github.com/tree-sitter/tree-sitter-python v0.23.6 h1:qHnWFR5WhtMQpxBZRwiaU5Hk/29vGju6CVtmvu5Haas= +github.com/tree-sitter/tree-sitter-python v0.23.6/go.mod h1:cpdthSy/Yoa28aJFBscFHlGiU+cnSiSh1kuDVtI8YeM= github.com/tree-sitter/tree-sitter-ruby v0.21.1-0.20240818211811-7dbc1e2d0e2d h1:fcYCvoXdcP1uRQYXqJHRy6Hec+uKScQdKVtMwK9JeCI= github.com/tree-sitter/tree-sitter-ruby v0.21.1-0.20240818211811-7dbc1e2d0e2d/go.mod h1:T1nShQ4v5AJtozZ8YyAS4uzUtDAJj/iv4YfwXSbUHzg= github.com/tree-sitter/tree-sitter-rust v0.21.3-0.20240818005432-2b43eafe6447 h1:o9alBu1J/WjrcTKEthYtXmdkDc5OVXD+PqlvnEZ0Lzc=