Analysing a Huge Codebase with Python

Question

I've written a Python program to analyse a huge C++ code base in excess of millions of lines of code. The job of the program is simply to search for all C and C++ style comments and extract all the comments which contain specific keywords and phrases and to write those comments to an associated file. So far the program works well enough. It's fairly quick and easy to configure and it does the job, but it doesn't search anywhere near as fast as I would like it to and I would like some advice on how to make things run faster.

Configuration

The program is very quick and easy to configure. If you'd like to search a codebase for a single expression all you have to do is pass it that expression as an argument and the directory to search in and it will do the rest for you. To search for comments containing the word "hack" in the src/ directory you could simply write:

./main.py -e hack -s ./src/

If you want to pass many expressions however, you need to use a specially crafted JSON file and pass the name of that file as an argument. An example JSON file might look like this:

[
  {
    "hack-stuff": [
      "hack",
      "hacked",
      "hacky"
    ]
  },
  "WARNING",
  "DANGER",
  [
    "bad",
    "badly"
  ]
]

The file is then passed to the program like such:

./main.py -f test-words.json -s ./src/

This will create four files: "hack-stuff", "WARNING", "DANGER" and "bad". The file "hack-stuff" will be where all comments containing the words hack, hacked and hacky will be placed, "bad" will be where all comments containing "bad" and "badly" will be placed, and "WARNING" and "DANGER" will be where comments containing "WARNING" and "DANGER" will be placed respectively. This example demonstrates the three ways (string, list or dictionary) you can specify which comments matches you want to place in which files.

Program Structure

At the core of the program is the SourceWalker class which contains all the internal logic required to analyse the codebase and write the comments to their respective files. There is also a main function in a separate file which reads the arguments into their relevant variables and performs the JSON processing (if necessary) before then initialising and running an instance of the class via the walk() method.

Performance

I've tried a number of things to make the program as performant as possible, including incorporating multiprocessing which yielded huge improvements, but I'm not sure what I can do to make it any faster at this point. The main slowdown is caused by the for loop within _process_files() on line 117. This is the part of the program which runs in each of the child processes and searches through each file looking for valid comments before checking them against a series of pre-compiled regular expressions to see if they match one of the expressions we're looking for. I'm sure there are better ways of extracting the comments from each file and then searching through them but I'm not sure what they would be. Any suggestions here would be greatly appreciated.

Additional Comments

While performance is my main concern here, I'd also appreciate any feedback on the correctness and style of my program. It seems to work as intended but I can't guarantee there aren't some edge cases I've missed during my testing.

The Code

main.py

#!/usr/bin/python3

import sys
import json
import os
import argparse
import SourceWalker

def initialiseParser():
    parser = argparse.ArgumentParser(description = "Search the contents of comments within source code files")

    parser.add_argument("--file_name", "--file", "-f", help = "Name of the file containing the JSON list of expressions to search for", type = str)
    parser.add_argument("--source-dir", "--src", "-s", help = "The root directory of the source files to search over", type = str)
    parser.add_argument("--output-dir", "--out", "-o", help = "The directory the output files will be placed in", type = str)
    parser.add_argument("--expression", "--expr", "-e", help = "The expression to search for within the source comments", type = str)
    parser.add_argument("--language", "--lang", "-l", help = "The style of comments to look for within the file", type = str)

    return parser

def main():
    parser = initialiseParser()
    args = parser.parse_args()

    if args.source_dir:
        source_dir = args.source_dir
    else:
        sys.exit("Source directory must be specified!")

    if args.file_name:
        file_name = args.file_name
        input_file = open(file_name, "r")
        expressions = json.loads(input_file.read())
    elif args.expression:
        expressions = []
        expressions.append(str(args.expression))
    else:
        sys.exit("Error: Expression or file containing expressions must be specified!")

    output_dir = "./comments/"

    if args.output_dir:
        output_dir = args.output_dir

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    walker = SourceWalker.SourceWalker(source_dir, output_dir, expressions, extensions=[ ".c", ".cpp", ".h", ".cxx" ])
    walker.walk()

    if "input_file" in locals():
        input_file.close()

    return

if __name__=="__main__":
    main()

SourceWalker.py #!/usr/bin/python3

import sys
import json
import re
import os
import multiprocessing
import codecs
from pathlib import Path

class SourceWalkerException(Exception):
    """
    Raised when there is an error processing the given expressions
    TODO: Make error handling more informative and clean up. Should display a message to the user explaing what went wrong and close all open files.
    """
    pass

class SourceWalker:
    _output_file_names = []
    _regexes = {}
    _expr_file_names = {}

    def __init__(self, source_dir, output_dir, expr_list, encoding = "ISO-8859-1", process_count = 12, extensions = [ ".c", ".h" ]):
        try:
            if not os.path.exists(source_dir) or not os.path.exists(output_dir):
                raise NotADirectoryError

            if process_count < 1:
                raise SourceWalkerException("Process count cannot be less than one!")

            codecs.lookup(encoding)

            if not isinstance(extensions, list):
                raise SourceWalkerException("Extensions must be passed as a list!")

            for extension in extensions:
                if extension[0] != '.':
                    raise SourceWalkerException("Extensions must start with a \'.\'!")
                elif len(extension) <= 1:
                    raise SourceWalkerException("Extensions must be more than one character long!")
        except NotADirectoryError as exception:
            raise SourceWalkerException("Directory does not exist! " + str(exception))
        else:
            self._source_dir = source_dir
            self._output_dir = output_dir
            self._encoding = encoding
            self._expr_list = expr_list
            self._process_count = process_count
            self._extensions = extensions
            self._process_expr_list()

    def _process_expr_list(self):
        for expr in self._expr_list:
            try:
                if isinstance(expr, list):
                    if len(expr) == 0:
                        raise SourceWalkerException("Expression list cannot be empty!")
                    
                    output_file_name = expr[0]

                    if not isinstance(output_file_name, str):
                        raise SourceWalkerException("Expression sub-lists can only contain strings!")
                    
                    for sub_expr in expr:
                        if not isinstance(sub_expr, str):
                            raise SourceWalkerException("Expression sub-lists can only contain strings!")
                        elif sub_expr in self._regexes.keys():
                            raise SourceWalkerException("Expressions can only appear once in the expression list!")

                        self._regexes[sub_expr] = re.compile("\s+%s(\s|,|:|;|\n)+" % (sub_expr)) # Naieve regex to catch expressions
                        self._expr_file_names[sub_expr] = self._output_dir + output_file_name

                    self._output_file_names.append(self._output_dir + output_file_name)
                elif isinstance(expr, dict):
                    if len(expr.keys()) == 0:
                        raise SourceWalkerException("Expression dictionary cannot be empty!")

                    output_file_name = list(expr)[0]

                    if not isinstance(expr[output_file_name], list):
                        raise SourceWalkerException("Expression dictionary cannot be empty!")
                    
                    for sub_expr in expr[output_file_name]:
                        if not isinstance(sub_expr, str):
                            raise SourceWalkerException("Expression sub-lists can only contain strings!")
                        elif sub_expr in self._regexes.keys():
                            raise SourceWalkerException("Expressions can only appear once in the expression list!")

                        self._regexes[sub_expr] = re.compile("\s+%s(\s|,|:|;|\n)+" % (sub_expr))
                        self._expr_file_names[sub_expr] = self._output_dir + output_file_name

                    self._output_file_names.append(self._output_dir + output_file_name)
                elif isinstance(expr, str):
                    if expr in self._regexes.keys():
                        raise SourceWalkerException("Expressions can only appear once in the expression list!")

                    self._output_file_names.append(self._output_dir + expr)
                    self._regexes[expr] = re.compile("\s+%s(\s|,|:|;|\n)+" % (expr))
                    self._expr_file_names[expr] = self._output_dir + expr
                else:
                    raise SourceWalkerException("Expression list can only contain dictionaries, lists, and strings!")
            except SourceWalkerException as exception:
                self.cleanup()
                raise

    def _process_files(self, input_files, output_files, mutexes): # Find way to process different types of source file, I'd rather not be limited to C only...
        for file_name in iter(input_files.get, None):
            with open(file_name, "r", encoding = self._encoding) as file_object:
                in_multi_comment = False
                in_single_comment = False
                in_string = False
                prev_char = ''
                comment = ''

                for line_num, line in enumerate(file_object, 1):
                    for char in line:
                        if char == '/':
                            if in_string or in_single_comment:
                                prev_char = char
                                continue
                            if prev_char == '*':
                                in_multi_comment = False
                                comment += char

                                for expr in self._regexes.keys():
                                    if self._regexes[expr].search(comment):
                                        mutexes[expr].acquire()
                                        os.write(output_files[expr], ("%s: %s %s\n" % (file_name, str(line_num), comment)).encode())
                                        mutexes[expr].release()

                                comment = ''
                            elif prev_char == '/':
                                in_single_comment = True
                                comment += prev_char
                        elif char == '*':
                            if in_string or in_single_comment or in_multi_comment:
                                if in_single_comment or in_multi_comment:
                                    comment += char

                                prev_char = char
                                continue
                            if prev_char == '/':
                                in_multi_comment = True
                                comment += prev_char
                        elif char == '"':
                            if prev_char == '\\' or in_single_comment or in_multi_comment:
                                prev_char = char
                                continue
                            
                            in_string = not in_string

                        prev_char = char

                        if in_single_comment or in_multi_comment:
                            comment += char

                    if in_single_comment:
                        in_single_comment = False
                        
                        for expr in self._regexes.keys():
                            if self._regexes[expr].search(comment):
                                mutexes[expr].acquire()
                                os.write(output_files[expr], ("%s: %s %s" % (file_name, str(line_num), comment)).encode())
                                mutexes[expr].release()

                        comment = ''

    def walk(self):
        input_files = multiprocessing.Queue(0)
        processes = []
        mutexes = {}
        output_files = {}

        for fname in self._output_file_names:
            try:
                file_handle = os.open(fname, os.O_WRONLY | os.O_CREAT)
                mutex = multiprocessing.Lock()
            except IOError:
                for file in output_files.keys():
                    output_files[file].close()

                raise SourceWalkerException("Error: Could not open output file %s, skipping!" % fname)

            for expr in self._expr_file_names.keys():
                if self._expr_file_names[expr] == fname:
                    output_files[expr] = file_handle
                    mutexes[expr] = mutex

        for root, dirs, file_names in os.walk(self._source_dir):
            for file_name in file_names:
                if any(ext in Path(file_name).suffix for ext in self._extensions):
                    input_files.put(os.path.join(root, file_name))

        for i in range(self._process_count):
            input_files.put(None)
        
        for cur_process in range(self._process_count):
            process = multiprocessing.Process(target = self._process_files, args = (input_files, output_files, mutexes))
            processes.append(process)
            process.start()

        for i in range(1, self._process_count):
            processes[i].join()

        for file in output_files.keys(): # Close the file associated with each expression
            try:
                os.close(output_files[file]) # Since multiple expressions can be associated with the same file we need to avoid invalid file closures
            except:
                pass

Answer 1

If I understand your description, you're only looking for comments but you are searching through the full code base every time. Since comments are normally a small part of the code (less than 10%?) I suggest doing a pre-process step first where you simply extract all the comments and then do the actual search on those.

By "extract" I mean save the comment in separate files so that you can search only in those files.

For repeated searches in the same codebase, this should help since the preprocessing can be done once and then the actual search has less text to look through.

Answer 2

You can specify required arguments in argparse rather than handling them yourself. You can also specify defaults, for example for output_dir.

os.makedirs takes exist_ok=True to indicate that it should only create the directory if it doesn't exist yet.

The extensions argument to the SourceWalker constructor probably shouldn't be hardcoded.

Use with open(file_name, "r") as file_handle: to make sure the file is always closed when leaving that context. On a related note, locals() should IMO only be used as a last resort, since it's hard to follow code where strings are used to refer to variables and IDEs can't do anything useful with that code. For example, if you were to rename input_file to input_file_path (I'd recommend that in any case) it would be trivial to forget to change the string reference.

I would recommend using black and isort to format the code. It'll be closer to idiomatic style that way, with no manual work.

After formatting I would recommend running flake8 to find other non-idiomatic code. For example it looks like the SourceWalker code is broken - the methods are at the same level as the class definition.

Creating an exception with a static string, such as SourceWalkerException("Process count cannot be less than one!"), is a code smell. It's not necessarily bad, but like boolean parameters it usually means something should be pulled apart. In this case it would be better to have a separate exception class for a too low process count (and the other error conditions), which could be something as simple as ProcessCountTooLowError. That way the only thing passed to an exception constructor (if anything) is whatever dynamic content can be used to debug that exception. Also, this means that if you ever end up wanting to handle the exception you can except ProcessCountTooLowError rather than having to parse the error message inside a generic except SourceWalkerException. (Oh, and custom exception classes should end in Error, not Exception.)

Rather than checking for things like whether the output directory exists I would let the application fail once it reaches the code which tries to write to that directory. This is called "time of check to time of use" - basically, whether the directory exists when you make that check has no bearing on whether it exists when the code actually tries to use it. It is also too narrow a check, since, for example, the directory could also not be writable by the current user, or the filesystem could be full.

On a similar note, running codecs.lookup(encoding) to check the encoding exists before actually using it in a completely different call seems like it could easily be an incomplete check. It might be better to constrain the encoding parameter options to only the available encodings. That way it's checked as early as possible, the users get a nice actionable error and the help text shows the possible values.

The plural of "regex" is "regexes".

Don't use mutable default arguments (discussion). On a related note, when should you use default arguments at all? Two rules of thumb:

1. If the default is not actually ever used because all the calls specify a value, the default is pointless.
2. If none of the calls override the default it is also pointless - it might as well be a variable or constant.

if not isinstance(extensions, list): is not idiomatic; it should be perfectly valid to pass in any iterable such as a tuple.

Constructors should, in general, do nothing more complex than setting field values. Once that's done a run, process or other method should be run separately to do the actual processing. I don't remember where I first saw this explained clearly, but see for example Why is using side effects bad practice in JavaScript constructors? and Constructors Must Be Code-Free

Answer 3

List literals

    expressions = []
    expressions.append(str(args.expression))

should just be

    expressions = [str(args.expression)]

Pathlib

This:

if not os.path.exists(output_dir):
    os.makedirs(output_dir)

should use the shiny new pathlib equivalent:

Path(output_dir).mkdir(exist_ok=True)

The same module can be used for

if not os.path.exists(source_dir) or not os.path.exists(output_dir):

Variable existence

if "input_file" in locals():
    input_file.close()

is sketchy. Usually the way to indicate in Python that a variable has a value or not is to potentially take None, not to potentially be undeclared. You can even mark it as maybe-none using the Optional type hint.

Return

The single return at the end of main() is redundant.

Dictionary length

if len(expr.keys()) == 0:

can be

if len(expr) == 0:

or even

if not expr:

Regex compilation

                self._regexs[sub_expr] = re.compile("\s+%s(\s|,|:|;|\n)+" % (sub_expr))

needs, at the least, a leading r to make that string literal raw.

Generally, taking input from a file and treating it as a non-validated sub-regular-expression is a bad idea. It's not eval-level bad, but it's not good. Are these actual regular expressions, or just substrings? If they're only substrings, call escape on them before inserting them into your outer regular expression.

If they actually are their own regular expression, you'll want to at least put each one in its own non-capturing group in the outer expression, to avoid nasty surprises.

Chained exceptions

except NotADirectoryError as exception:
    raise SourceWalkerException("Directory does not exist! " + str(exception))

should be

except NotADirectoryError as exception:
    raise SourceWalkerException("Directory does not exist! " + str(exception)) from exception