Counting density of Python code

Question

I created a tool to examine Python 3 source code. It measures the density of program. More info in docstring (to not reinvent the wheel).

What could be done to increase its usability? Any edge cases that I didn't think of?

#!/usr/bin/env python3

'''
Module counts how dense is Python source code in
terms of lexical entities per LoC.
'''

import ast


def traverse(node, line_nos):
    '''
    go through the AST and count lines
    '''
    children = list(ast.iter_child_nodes(node))

    try:
        current_line = node.lineno
    except AttributeError:
        current_line = 0

    count = 1
    line_nos.update([current_line])

    if len(children) == 0:
        return 1

    for child in ast.iter_child_nodes(node):
        count += traverse(child, line_nos)

    return count


def count_code_entities(source):
    '''
    parse the source&count
    '''
    line_nos = set()
    tree = ast.parse(source)

    # len - 1 to compensate for default value 0
    # which cannot hold any code 
    return traverse(tree, line_nos), len(line_nos) - 1


def main():
    '''
    parsin&stuff
    '''
    from argparse import ArgumentParser, FileType
    parser = ArgumentParser()
    parser.add_argument('infile',
                        help='source file to analyze',
                        type=FileType(),
                        default='-',
                        nargs='?',)

    args = parser.parse_args()

    with args.infile as file:
        source = file.read()
        name = file.name

    count, number = count_code_entities(source)
    print(name, ': ', round(count/number, 4), 'entities/line of code')


if __name__ == '__main__':
    main()

Answer 1

1. The docstrings are vague. Documentation needs to be precise and complete. It should answer questions like "what arguments should I pass?" and "what does it return?" For example, "go through the AST and count lines" should become something like, "Return the number of nodes in the AST rooted at node, and update the set line_nos with the line numbers of these nodes, plus the line number 0 if any node has no line number."

2. It's not clear that the "plus the line number 0 if any node has no line number" behaviour is needed. Instead of doing this and then having to subtract one later to compensate, why not avoid adding the line number 0 in the first place?

3. When traversing an AST, it's often convenient to use the ast.NodeVisitor class, like this:

   class DensityVisitor(ast.NodeVisitor):
       """AST visitor that accumulates the count of nodes and the set of line
       numbers in the AST.
       """
       def __init__(self):
           self.node_count = 0
           self.line_numbers = set()
   
       def visit(self, node):
           self.node_count += 1
           try:
               self.line_numbers.add(node.lineno)
           except AttributeError:
               pass
           self.generic_visit(node)
   
       @property
       def density(self):
           """The density of code (nodes per line) in the visited AST."""
           return self.node_count / len(self.line_numbers)
   

   and then in main you'd write:

   visitor = DensityVisitor()
   visitor.visit(ast.parse(source))
   print('density (nodes/line) =', round(visitor.density, 4)))
   

   Update: Why would you prefer ast.NodeVisitor over explicitly walking the tree using ast.iter_child_nodes plus recursion? Well, because of separation of concerns — the logic for walking the tree is delegated to the ast.NodeVisitor base class, and you only have to specify what happens at each node. (Another possibility would be to use ast.walk.)

Answer 2

Great job overall, the code looks quite clean.

There are only few things that I would mention:

• you can calculate count using sum() one-liner
• if len(children) == 0: can be replaced with if not children:
• and, we don't actually need to handle this base case, since we are going to use sum with +1, and in case of children being an empty list, we'll have 1 as a result naturally
• docstrings need to be in triple double-quotes, start with a capital letter and end with a dot
• move import statements to the top of the module
• file is not a good variable name since it shadows the built-in file keyword
• let's rename main function to a more readable report_code_density
• use a formatted string to print out the report - it will be a bit better in terms of having the report printed consistently on both Python2.x and Python3.x
• you can use getattr() instead of try/except for the lineno (credits to @peilonrayz)

Improved code:

#!/usr/bin/env python3

"""Module counts how dense is Python source code in terms of lexical entities per LoC."""

from argparse import ArgumentParser, FileType
import ast


def traverse(node, line_nos):
    """Goes through the AST and count lines."""
    current_line = getattr(node, 'lineno', 0)
    line_nos.update([current_line])

    return sum(traverse(child, line_nos) for child in ast.iter_child_nodes(node)) + 1


def count_code_entities(source):
    """Parses the source & counts."""
    line_nos = set()
    tree = ast.parse(source)

    # len - 1 to compensate for default value 0, which cannot hold any code
    return traverse(tree, line_nos), len(line_nos) - 1


def report_code_density():
    """Main execution block of the program."""
    parser = ArgumentParser()
    parser.add_argument('infile',
                        help='Source file to analyze',
                        type=FileType(),
                        default='-',
                        nargs='?')

    args = parser.parse_args()

    with args.infile as source_file:
        source = source_file.read()
        name = source_file.name

    count, number = count_code_entities(source)
    print('{name}: {value} entities/line of code'.format(name=name, value=round(count / number, 4)))


if __name__ == '__main__':
    report_code_density()

And, here is the pylint report (with default settings) for the improved code:

$ pylint test.py
Report
======
24 statements analysed.

Global evaluation
-----------------
Your code has been rated at 10.00/10