From 98f1d8097885fe0927aa4a48f72f91d847dd7e1c Mon Sep 17 00:00:00 2001
From: Alex Torok <alext9@gmail.com>
Date: Mon, 26 Aug 2024 20:19:59 +0000
Subject: [PATCH] Sample implementation of lazy indexing

---
 gazelle/python/generate.go | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/gazelle/python/generate.go b/gazelle/python/generate.go
index c563b47bf3..fa3d0fccf0 100644
--- a/gazelle/python/generate.go
+++ b/gazelle/python/generate.go
@@ -487,6 +487,35 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 		os.Exit(1)
 	}
 
+	dirsToIndex := map[string]struct{}{}
+	for _, r := range result.Gen {
+		modulesRaw := r.PrivateAttr(config.GazelleImportsKey)
+		modules := modulesRaw.(*treeset.Set)
+		it := modules.Iterator()
+		for it.Next() {
+			mod := it.Value().(module)
+			currentDir := ""
+			parts := strings.Split(mod.Name, ".")
+			for i := range parts {
+				if i == 0 {
+					currentDir = parts[i]
+				} else {
+					currentDir = currentDir + "/" + parts[i]
+				}
+				dirsToIndex[currentDir] = struct{}{}
+			}
+		}
+	}
+
+	dirsToIndex[""] = struct{}{}
+
+	toIndex := make([]string, 0, len(dirsToIndex))
+	for dir := range dirsToIndex {
+		toIndex = append(toIndex, dir)
+	}
+
+	result.PackagesToIndex = toIndex
+
 	return result
 }