metrieke/python.go

package metrieke

import (
	"fmt"
	"io/ioutil"
	"strings"

	sitter "github.com/smacker/go-tree-sitter"
	"github.com/smacker/go-tree-sitter/python"
)

type PythonFile struct {
	Source  []byte
	Path    string
	Root    *sitter.Node
	Metrics *PythonFileMetrics
}

type PythonFileMetrics struct {
	Loc                int
	Lloc               int
	Dloc               int
	Cloc               int
	DocumentationLines int
	AvgComplexity      float64
	SumComplexity      int
	AvgNesting         float64
	SumNesting         int
}

func NewPythonFileFromContents(path string, contents []byte) *PythonFile {
	parser := sitter.NewParser()
	parser.SetLanguage(python.GetLanguage())
	tree := parser.Parse(nil, contents)
	node := tree.RootNode()

	return &PythonFile{
		Path:    path,
		Source:  contents,
		Root:    node,
		Metrics: &PythonFileMetrics{},
	}
}

func NewPythonFile(path string) *PythonFile {
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		fmt.Println("got error")
	}
	return NewPythonFileFromContents(path, contents)
}

func (pf *PythonFile) Parse() {
	pf.pythonLines()

	//fmt.Printf(string(pf.Source))
	numMethods := 0
	sumComplexity := 0
	var sumNesting int
	var avgNesting float64

	mccc := make([]int, 0)

	pattern := []byte(`(function_definition name: (identifier) @function-name)`)
	queryChildNodes(pattern, python.GetLanguage(), pf.Root, func(node *sitter.Node) {
		//fmt.Printf("node: %v\n", node.Parent())
		complexity := pythonMcCabeComplexity(node.Parent())
		sumNesting += pythonMaxNesting(node.Parent())
		numMethods += 1
		sumComplexity += complexity
		mccc = append(mccc, complexity)
	})

	avgComplexity := float64(0)
	if numMethods > 0 {
		avgComplexity = float64(sumComplexity) / float64(numMethods)
		avgNesting = float64(sumNesting) / float64(numMethods)
	}
	pf.Metrics.AvgComplexity = avgComplexity
	pf.Metrics.SumComplexity = sumComplexity
	pf.Metrics.SumNesting = sumNesting
	pf.Metrics.AvgNesting = avgNesting
}

// handle the simple line stuff parsing
func (pf *PythonFile) pythonLines() {
	lines := strings.Split(string(pf.Source), "\n")
	loc := len(lines)
	lloc := 0
	dloc := 0
	cloc := 0
	inComment := false
	for _, line := range lines {
		trimmedLine := strings.TrimSpace(line)
		if inComment == false && strings.HasPrefix(trimmedLine, `"""`) {
			inComment = true
			trimmedLine = strings.Replace(trimmedLine, `"""`, "", 1)
		}

		if inComment {
			dloc += 1
		}

		if inComment && strings.HasSuffix(trimmedLine, `"""`) {
			inComment = false
			continue
		}

		if strings.HasPrefix(trimmedLine, "#") {
			cloc += 1
			continue
		}

		if !inComment && trimmedLine != "" {
			lloc += 1
		}
	}

	pf.Metrics.Loc = loc
	pf.Metrics.Lloc = lloc
	pf.Metrics.Dloc = dloc
	pf.Metrics.Cloc = cloc
}

func PythonMethodDocumentation(node *sitter.Node) int {
	numMethods := 0
	pattern := []byte(`(function_definition(block(expression_statement(string) @functiondoc)))`)
	queryChildNodes(pattern, python.GetLanguage(), node, func(node *sitter.Node) {
		fmt.Printf("doc %v\n", node)
	})

	return numMethods
}

func PythonNumMethods(node *sitter.Node) int {
	numMethods := 0
	pattern := []byte(`(function_definition name: (identifier) @function-name)`)
	queryChildNodes(pattern, python.GetLanguage(), node, func(node *sitter.Node) {
		numMethods += 1
		//fmt.Printf("%v\n", node.Parent())
	})

	return numMethods
}

// python method level metrics
func pythonMaxNesting(node *sitter.Node) int {
	maxChild := -1
	//fmt.Printf("node\n")
	for i := 0; i < int(node.NamedChildCount()); i++ {
		childDepth := pythonMaxNesting(node.NamedChild(i))
		//fmt.Printf("dept %v of child %v (%v)\n", childDepth, i, node.NamedChild(i).Type())
		if childDepth >= maxChild {
			maxChild = childDepth
		}
	}

	//fmt.Printf("return max dept %v+1\n", maxChild)
	if node.Type() == "block" {
		//fmt.Printf("block\n")
		return maxChild + 1
	} else {
		return maxChild
	}
}

// traverse down the method node to get the cyclomatic complexity
func pythonMcCabeComplexity(node *sitter.Node) int {
	iter := sitter.NewNamedIterator(node, sitter.BFSMode)
	complexity := 1 // 1 per method
	for {
		n, err := iter.Next()
		if err != nil {
			return complexity
		}
		// +1 for every branch statement
		if n.Type() == "if_statement" {
			complexity += 1
		}
		// seems this does not exist for python in tree-sitter (or these go-bindings)
		if n.Type() == "elif_statement" {
			complexity += 1
		}
		if n.Type() == "for_statement" {
			complexity += 1
		}
		if n.Type() == "while_statement" {
			complexity += 1
		}
		if n.Type() == "with_statement" {
			complexity += 1
		}
		if n.Type() == "except_clause" {
			complexity += 1
		}
		// this is something radon does, I am not sure if we should add this
		//if n.Type() == "boolean_operator" {
		//complexity += 1
		//}
	}
}