Python: code a Flake8 plugin to check your own rule

Conventions is not only a code style guide. If you include some rules to ensure security then you add some work for the reviewer. It may be interesting to automate this check with some code.

Flake8 is a tool to ensure your code follow some rules. You can add plugins and it's easy to create your own. This post is the textual version of this very good Youtube tutorial by anthonywritescode.

Non-syntax conventions

You can use pylint or Black to reformat and uniformize your code but conventions may contains some other rules:

• Security: your company may ban some function like os.popen() or forbid a call to request.get() without the cert attribute.

• Log message: you may have some rule to have explicit and descriptive log entry, for example logger.critical() must contains some specific keywords (error uid, tracing span...).

Automate your conventions checks

Building a flake8 has two major advantages:

• the review checklist is getting thinner because the continious integration tool take this job.
• the developer can be alerted by this noncompliance and fix-it himself instead of waiting pair-review.

If you want to forbid the add of a controler class without documentation, you only have to ensure the docsting is set: It's easier for the reviewer to detect empty docstring than looking at each controler class if they contains documentation.

Automate the first part of the checks (the existence) to help the reviewer to focus on validating the content. So some rules can only be partially automated but it still reduce the mental load of the reviewer.

Follow your technical debt

If you have many projects and you decide to migrate from library A to library B you can use flake8 to alert you about remaining library A usage.

Once your internal_conventions plugin is set and running it's easy to add some rules with some new error-code.

Create Flake8 plugin

It's very easy to check your own rules with a Flake8 plugin (not only because we don't have to use regex), the bootstrap is small.

In this example I write a plugin to check every call to logger.info() contains mandatory_arg. So logger.info("abc", mandatory_arg=1) is valid while logger.info("abc") is invalid.

My project name is flake8_check_logger, flake8_ is a prefix for flake8 plugin (this good naming convention help us to discover plugin on Github). We use the error code LOG042 with message "Invalid logger". By the way use a prefix not already used by your flake8 plugins.

Setup

setup.py:

from setuptools import setup
setup()

setup.cfg:

[metadata]
name = flake8_check_logger
version = 0.1.0

[options]
py_modules = flake8_check_logger
install_requires = 
    flake8>=3.7

[options.entry_points]
flake8.extension =  
    LOG=flake8_check_logger:Plugin

flake8_check_logger.py

import importlib.metadata
import ast
from typing import Any, Generator, Type, Tuple


class Visitor(ast.NodeVisitor):
    def __init__(self):
        self.problems = []

    def visit_Call(self, node: ast.Call) -> None:
        # TODO we write our test here

        self.generic_visit(node)


class Plugin:
    name = __name__
    version = importlib.metadata.version(__name__)

    def __init__(self, tree: ast.AST):
        self._tree = tree

    def run(self) -> Generator[Tuple[int, int, str, Type[Any]], None, None]:
        visitor = Visitor()
        visitor.visit(self._tree)
        for line, col, code, message in visitor.problems:
            yield line, col, f"LOG{code} {message}", type(self)

 Discover AST (and astpretty)

Instead of using regex we will use Abstract Syntax Trees. AST transform any python source code to an object representation.

We will install astpretty with pip for development purpose only, it's not a requirement to add in our plugin.

Then you can use it on your terminal:

astpretty /dev/stdin <<< "logger.info('msg', mandatory_arg=1)"

And we obtain some information like

Expr(
            value=Call(
                func=Attribute(
                    value=Name(id='logger', ctx=Load()),
                    attr='info',
                    ctx=Load(),
                ),
                args=[Constant(value='msg', kind=None)],
                keywords=[
                    keyword(
                        arg='mandatory_arg',
                        value=Constant(value=1, kind=None),
                    ),
                ],
            ),
        ),

Write the test

from flake8_check_logger import Plugin

import ast
from typing import Set


def _results(source: str) -> Set[str]:
    tree = ast.parse(source)
    plugin = Plugin(tree)
    return {f"{line}:{col} {msg}" for line, col, msg, _ in plugin.run()}

def test_valid_logger_call():
    assert _results("logger.info('msg', mandatory_arg=1)") == set()

def test_invalid_logger_call_no_mandatory_arg():
    ret = _results("logger.info('msg')")
    assert ret == {'1:0 LOG042 Invalid Logger'}

You can add more testcases: logger.something_else() or a local function with similar name info() should not require this argument, test an empty line, etc.

Write the check

I split my test in two parts for lisibility, this code could be refactored but I prefer to use very basic syntax for this tutorial.

def visit_Call(self, node: ast.Call) -> None:
        mandatory_arg_found = False
        is_right_function = False

        if hasattr(node.func, "value") and node.func.value.id == "logger":
            if node.func.attr == "info":
                is_right_function = True

        if is_right_function:
            for keyword in node.keywords:
                if keyword.arg == "mandatory_arg" and keyword.value.value != None:
                    mandatory_arg_found = True

            if not mandatory_arg_found:
                self.problems.append((
                    node.lineno,
                    node.col_offset, 
                    042,
                    "Invalid logger"
                ))

        self.generic_visit(node)

Local test

We install our plugin with:

pip install -e .

.

And you can check flake load this plugin with

flake8 --version

.

Go further

Your can find more information on the flake8 documentation with a link to the very helpful video tutorial (30 minutes).

With python you can access to something with code: __doc__ for documentation or __bases__ for parent class. Theses dunder function can help you to write some specific rules.

You can learn more about existing flake8 plugin thanks to Jonathan Bowman article and read the source-code of theses plugins to help you if you are blocked by a complex case.