github.com/ismailbayram/bigpicture@v0.0.0-20231225173155-e4b21f5efcff/internal/browser/python.go (about) 1 package browser 2 3 import ( 4 "fmt" 5 "github.com/ismailbayram/bigpicture/internal/graph" 6 "os" 7 "regexp" 8 "strings" 9 ) 10 11 type PythonBrowser struct { 12 ignoredPaths []string 13 moduleName string 14 rootDir string 15 tree *graph.Tree 16 } 17 18 func (b *PythonBrowser) Browse(parentPath string) { 19 b.moduleName = b.getModuleName() 20 b.browse(parentPath, b.tree.Root) 21 b.clearNonProjectImports() 22 } 23 24 func (b *PythonBrowser) getModuleName() string { 25 directory, err := os.Getwd() 26 if err != nil { 27 panic(err) 28 } 29 return strings.Split(directory, "/")[len(strings.Split(directory, "/"))-1] 30 } 31 32 func (b *PythonBrowser) clearNonProjectImports() { 33 for _, node := range b.tree.Nodes { 34 var clearedImports []string 35 for _, imp := range node.ImportRaw { 36 if _, ok := b.tree.Nodes[imp]; ok { 37 clearedImports = append(clearedImports, imp) 38 } 39 } 40 node.ImportRaw = clearedImports 41 } 42 } 43 44 func (b *PythonBrowser) browse(parentPath string, parentNode *graph.Node) { 45 entries, err := os.ReadDir(parentPath) 46 47 if err != nil { 48 panic(err) 49 } 50 51 for _, e := range entries { 52 fName := e.Name() 53 path := fmt.Sprintf("%s/%s", parentPath, fName) 54 if isIgnored(b.ignoredPaths, path) { 55 continue 56 } 57 58 if e.IsDir() && !strings.Contains(fName, ".") { 59 node := graph.NewNode(fName, path, graph.Dir, nil) 60 b.tree.Nodes[node.Path] = node 61 b.browse(path, node) 62 parentNode.LineCount += node.LineCount 63 } else if strings.HasSuffix(fName, ".py") { 64 node := b.parseFile(path, parentNode) 65 b.tree.Nodes[node.Path] = node 66 parentNode.LineCount += node.LineCount 67 } 68 } 69 } 70 71 func (b *PythonBrowser) parseFile(path string, parentNode *graph.Node) *graph.Node { 72 file, err := os.ReadFile(path) 73 if err != nil { 74 panic(err) 75 } 76 fileContent := string(file) 77 78 // extract functions 79 functions := make([]graph.Function, 0) 80 functionsInfo := b.findFunctions(fileContent) 81 for name, functionBody := range functionsInfo { 82 functions = append(functions, graph.Function{ 83 Name: name, 84 LineCount: len(strings.Split(functionBody, "\n")) - 1, 85 }) 86 } 87 88 fName := path[strings.LastIndex(path, "/")+1:] 89 node := graph.NewNode(fName, path, graph.File, b.findImports(fileContent)) 90 node.Functions = functions 91 node.LineCount = strings.Count(fileContent, "\n") 92 93 return node 94 } 95 96 func (b *PythonBrowser) findImports(pythonCode string) []string { 97 var imports []string 98 99 lines := strings.Split(pythonCode, "\n") 100 importRegex := regexp.MustCompile(`^\s*import\s+([^\s#]+)`) 101 fromImportRegex := regexp.MustCompile(`^\s*from\s+([^\s]+)\s+import`) 102 103 for _, line := range lines { 104 var importItem string 105 106 if matches := importRegex.FindStringSubmatch(line); len(matches) > 1 { 107 importItem = b.rootDir + strings.Replace(matches[1], ".", "/", -1) + ".py" 108 } else if matches := fromImportRegex.FindStringSubmatch(line); len(matches) > 1 { 109 importItem = b.rootDir + strings.Replace(matches[1], ".", "/", -1) + ".py" 110 } 111 if importItem == "" { 112 continue 113 } 114 115 if strings.HasSuffix(importItem, "*.py") { 116 importItem = importItem[:len(importItem)-5] 117 } 118 imports = append(imports, importItem) 119 } 120 121 return imports 122 } 123 124 func (b *PythonBrowser) findFunctions(fileContent string) map[string]string { 125 functions := make(map[string]string) 126 127 lines := strings.Split(fileContent, "\n") 128 functionRegex := regexp.MustCompile(`def\s+([a-zA-Z_]\w*)\s*\(([^)]*)\)\s*(?:->\s*([a-zA-Z_]\w*|\w*))?:`) 129 130 var currentFunctionName string 131 var currentFunctionContent string 132 133 for _, line := range lines { 134 if matches := functionRegex.FindStringSubmatch(line); len(matches) > 1 { 135 if currentFunctionName != "" { 136 functions[currentFunctionName] = strings.TrimSpace(currentFunctionContent) 137 currentFunctionContent = "" 138 } 139 140 currentFunctionName = matches[1] 141 } 142 143 if currentFunctionName != "" && (strings.Contains(line, "@") || strings.Contains(line, "class ")) { 144 functions[currentFunctionName] = strings.TrimSpace(currentFunctionContent) 145 currentFunctionContent = "" 146 currentFunctionName = "" 147 continue 148 } 149 150 if currentFunctionName != "" { 151 currentFunctionContent += line 152 currentFunctionContent += "\n" 153 } 154 } 155 156 if currentFunctionName != "" { 157 functions[currentFunctionName] = strings.TrimSpace(currentFunctionContent) 158 } 159 160 return functions 161 }