用Golang实现一个统计代码行数的程序

本文链接：https://blog.csdn.net/Peerless__/article/details/121191005

在开发中，我们有时候想要统计一个文件夹下的所有代码的行数，但是有时候找这样的工具也挺麻烦的。那么就自己实现一个吧。

思路：
1、通过命令行参数获取要统计的代码所在根目录以及代码文件的后缀，比如Go语言是.go C++语言是.cpp，可以同时统计多种类型的文件

2、从根目录开始递归遍历文件：1.如果是目录，就递归遍历；2.如果是普通文件，根据后缀判断是否是要统计的代码的文件，如果是就统计并打印代码行数。

代码如下：

package main

import (
	"bufio"
	"fmt"
	flag "github.com/spf13/pflag"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"sync"
)

var (
	suffixLines   map[string]int
	totalLines    int
	rootPath      string
	defaultSuffix string = ".go"
	suffix        map[string]struct{}

	wg           sync.WaitGroup
	staticticsWg sync.WaitGroup
)

type Line struct {
	Filepath string
	Lines    int
}

func main() {
	// 先获取要统计的代码所在的rootPath 以及代码文件名后缀
	err := getPath()
	if err != nil {
		fmt.Println("Get Current Path error:", err)
		return
	}
	// 如果不是目录，直接统计文件行数，然后返回
	ok := IsDir(rootPath)
	if !ok {
		line := getFileLine(rootPath)
		fmt.Printf("%s: %d\n", rootPath, line)
		return
	}

	// 统计目录中所以指定后缀文件的行数
	staticticsWg.Add(1)
	staticsChan := make(chan *Line, 128)
	go statictics(staticsChan)

	// 遍历目录下的文件
	wg.Add(1)
	pathSeparator := string(os.PathSeparator)
	go staticticsAllFile(pathSeparator, rootPath, staticsChan)

	// wait返回表示目录已经遍历完毕
	wg.Wait()

	// 通知统计协程退出
	close(staticsChan)

	staticticsWg.Wait()

	for name, lines := range suffixLines {
		fmt.Printf("File suffix %s total line: %d\n", name, lines)
	}
	fmt.Println("Total Line: ", totalLines)
}

// 给rootPath以及后缀赋值
func getPath() error {

	dafaultFile, err := os.Getwd()
	if err != nil {
		return err
	}
	filename := flag.String("f", dafaultFile, "-f filePath")
	sfxs := flag.StringSlice("s", []string{defaultSuffix}, "-s suffix1,...,suffixn")

	flag.Parse()

	// 如果是否是相对路径
	if (*filename)[0] != '.' {
		rootPath = *filename
	} else {
		rootPath, _ = filepath.Abs(*filename)
	}

	suffix = make(map[string]struct{}, len(*sfxs))
	for i := 0; i < len(*sfxs); i++ {
		key := (*sfxs)[i]
		suffix[key] = struct{}{}
	}

	fmt.Printf(`Statistics target Root Path: "%s"`, rootPath)
	fmt.Println()
	fmt.Println(`File Suffix: `, *sfxs)
	fmt.Println()
	fmt.Println()
	fmt.Println("--------------------------------------------------------------------------------")
	fmt.Println()

	return nil
}

// 打印并统计代码总行数
func statictics(channel chan *Line) {
	defer staticticsWg.Done()

	suffixLines = make(map[string]int, len(suffix))

	for {
		line, ok := <-channel
		if !ok {
			return
		}
		totalLines += line.Lines
		suffixLines[path.Ext(line.Filepath)] += line.Lines
		fmt.Printf("%s: %d\n", line.Filepath, line.Lines)
	}

}

// 遍历所有文件，如果是普通文件就统计行数，如果是目录，递归遍历目录
func staticticsAllFile(pathSeparator string, fileDir string, channel chan *Line) {
	files, _ := ioutil.ReadDir(fileDir)

	for _, file := range files {
		filePath := fileDir + pathSeparator + file.Name()
		if file.IsDir() {
			if (file.Name())[0] == '.' { // 忽略隐藏文件
				continue
			}
			wg.Add(1)
			go staticticsAllFile(pathSeparator, filePath, channel)
		} else if _, ok := suffix[path.Ext(file.Name())]; ok {
			line := getFileLine(filePath)
			channel <- &Line{
				Filepath: filePath,
				Lines:    line,
			}

		}
	}

	wg.Done()
}

// 获取文件的行数
func getFileLine(filename string) (line int) {
	file, err := os.Open(filename)
	if err != nil {
		fmt.Println("getFileLine error, file name:", filename, " error:", err)
		return 0
	}
	defer file.Close()

	line = 0
	reader := bufio.NewReader(file)
	for {
		_, isPrefix, err := reader.ReadLine()
		if err != nil {
			break
		}
		if !isPrefix {
			line++
		}
	}

	return
}

// 判断所给路径是否为文件夹
func IsDir(path string) bool {
	s, err := os.Stat(path)
	if err != nil {
		return false
	}
	return s.IsDir()
}