Python: code a Flake8 plugin to check your own rule
Tags: python, automation, codequality
Published: 2023-03-04 11:00:00

---
title: Python: code a Flake8 plugin to check your own rule 
published: true
description: how to write a flake8 plugin to your personal conventions like checking some optional arguments are always passed to a specific function
series: tricks
tags: python,automation,quality
---

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](https://pypi.org/project/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](https://www.youtube.com/watch?v=ot5Z4KQPBL8) 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
```python
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](https://pypi.org/project/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

```python
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.

```python
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](https://flake8.pycqa.org/en/latest/plugin-development/index.html) 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](https://dev.to/bowmanjd/some-flake8-plugins-for-python-linting-107h) and read the source-code of theses plugins to help you if you are blocked by a complex case.