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  }