Lines in code demo for potential use un flake8 plugin

lines_in_code.py

import ast

func_code = """
def test():
    foo = 1
    bar = 2
    return foo + bar
"""

tree = ast.parse(func_code)
# in this case is just a single node of type ast.FunctionDef
func = tree.body[0]

def lines_in_func(node):
    """
    Count the lines of a FunctionDef node
    """
    # use if!, this is just for "show"
    assert isinstance(node, ast.FunctionDef)
    if node.body:
        return (node.body[-1].lineno - node.body[0].lineno) + 1
    else:
        return 1


print(func_code)
print("Lines in function:", lines_in_func(func))
assert lines_in_func(func) == 3

func_code_with_no_return = """
def test():
    foo = 1
    bar = 2
    foobar = 3
    print(foo + bar + foobar)
"""
print(func_code_with_no_return)
print("Lines in function: ",
      lines_in_func(ast.parse(func_code_with_no_return).body[0]))
assert lines_in_func(ast.parse(func_code_with_no_return).body[0]) == 4


func_code_weird_assignment = """
def test():
    foo = 1;  bar = 2; foobar = 3
    print(foo + bar + foobar)
"""
print(func_code_weird_assignment)
print("Lines in function: ",
      lines_in_func(ast.parse(func_code_weird_assignment).body[0]))
assert lines_in_func(ast.parse(func_code_weird_assignment).body[0]) == 2



func_code_weird_line_breaks = """
def test():
    foo = 1
    bar = 2
    foobar = 3
    return foo + \
             bar + \
               foobar
"""
print(func_code_weird_line_breaks)
print("Lines in function: ",
      lines_in_func(ast.parse(func_code_weird_line_breaks).body[0]))
assert lines_in_func(ast.parse(func_code_weird_line_breaks).body[0]) == 4


class_src = """
class Test:
    def foo(self):
        name = "foo"
        return name + "bar"
"""
klass = ast.parse(class_src)
foo_method = [
    node for node in ast.walk(klass)
    if isinstance(node, ast.FunctionDef)
].pop() # extract foo method definition from the tree

print(class_src)
print("Lines in foo method: ", lines_in_func(foo_method))
assert lines_in_func(foo_method) == 2