Flake8 Chronicles: Tale of crafting custom plugins

Hello, I’m Shreyansh 👋🏼, working in Fyle's engineering team. Over time our backend team grew in size, each person coding in their specific style which caused trouble in reviewing PRs. We decided to enforce some coding standards to ease reviews and have some mercy on ourselves 🫠

Enter flake8 the savior and its flexibility to create custom plugins for project-specific needs. Adopting flake8 has been very beneficial in maintaining a consistent coding style across our project. Wait a second, what is flake8 though? 🤔

“Flake8 is a Python tool that checks your Python codebase for errors, styling issues, and complexity.”

We use flake8 to run linting checks against our code. What are lint checks? These are basic coding guidelines and rules that ensure good code quality.

If you use Flake8 in your project already or plan to use it, this may be a good read for you. This will give you insights on:

• How to create a custom flake8 plugin 🚀

• How to publish your plugin 💥

• How to plan and adopt a plugin incrementally in an existing code base 📈

We at Fyle have built multiple plugins over time; I’ll talk about the ones that can be helpful in your projects.

1. Restrict Imports - Fine-tune Your Project's Dependency Landscape

   Plugin Link: Flake8 Restrict Imports

   When working on large projects, managing dependencies is a delicate balancing act. The Flake8 Restrict Imports plugin enables developers to enforce strict control over the imported modules and packages. This helps adhere to project-specific guidelines, prevent the accidental use of disallowed dependencies, and ultimately foster a more controlled and organized codebase and say goodbye to circular import issues as well 👋🏼

# flake8: noqa
# Allowed imports
import os
from my_module import allowed_function

# Disallowed imports
import insecure_module  # Flake8 will raise an error here

2. Enforce Type Annotations - Boost Code Readability and Reliability

   Plugin Link: Flake8 Enforce Type Annotations

   With the rise of static typing in Python, incorporating type annotations has become a best practice. This plugin ensures that functions and methods include proper type hints, promoting better code readability and reliability. This plugin contributes to the overall maintainability of the codebase by catching missing type annotations and encouraging adherence to the declared types.

# Missing type annotations will trigger Flake8 warning
def add_numbers(a, b):
    return a + b

# With type annotations
def add_numbers(a: int, b: int) -> int:
    return a + b

3. Enforce kwargs - Ensure Function Arguments Adherence

   Plugin Link: Flake8 Enforce kwargs

   Function arguments are the building blocks of any Python application, and consistent usage is key to code maintainability. The Enforce kwargs Flake8 plugin assists developers in maintaining uniformity by enforcing the use of keyword arguments in function definitions and function calls. This leads to clearer and more self-documenting functional code from the caller side 📞

# function defined without keyword arguments
def add_numbers(a, b):
    return a + b

# function defined with keyword arguments
def add_numbers(*, a, b):
		return a + b

# calling the function with keyword arguments
add_numbers(1, 2) # This is incorrect
add_numbers(a=1, b=2) # This is correct

Now let’s move on to the spicy bit of how to create your custom flake8 plugin 🌶️

Create a custom flake8 plugin 🚀

Let’s take the example of our enforce type annotations plugin to understand this process better.

Step 1: Setup project directory

mkdir flake8_enforce_type_annotations
cd flake8_enforce_type_annotations

Step 2: Create the plugin module. Let’s name it enforce_type_annotations.py

touch enforce_type_annotations.py 

Step 3: Add the plugin class.

class Plugin:
    name = "enforce_type_annotations"
    version = "0.1.0"

    def __init__(self, tree: ast.AST, filename: str) -> None:
        self.tree = tree
        self.current_filename = filename
        path = os.path.splitext(filename)[0]
        mod_path = []
        while path:
            if os.path.exists(os.path.join(path, '.flake8')):
                break
            dir, name = os.path.split(path)
            mod_path.insert(0, name)
            path = dir
        self.current_module = '.'.join(mod_path)

    def run(self) -> Generator:
        errors: List[Tuple[int, int, str]] = []
        visitor = FunctionVisitor(current_module=self.current_module, current_filename=self.current_filename, errors=errors)
        visitor.visit(self.tree)
        for lineno, colno, msg in errors:
            yield lineno, colno, msg, type(self)

This should accept the tree and filename in its constructor. tree argument is mandatory for the flake8 plugin to work since flake8 runs on our Python code’s ast tree.

What is ast tree, you ask? 😕

This is just another representation of our Python code in a tree format. Each node in the tree is an object defined in Python. Let me show an example to explain this jargon better.

For our python code:

print('Hello, world!')

The ast tree generated would be:

Module(
    body=[
        Expr(
            value=Call(
                func=Name(id='print', ctx=Load()),
                args=[
                    Constant(value='Hello, world!')],
                keywords=[]))],
    type_ignores=[])

Okay, ast is interesting, but how does this help with our flake8 plugin? 🤔

Flake8 uses ast module to parse and analyze the code into a tree, traverses all the nodes in the tree, and checks and reports issues.

Now, coming back to our plugin class. We defined a run method that yields a generator for the lint errors.

Our logic for the plugin will reside in FunctionVisitor, which inherits from ast.NodeVisitor, which gets called for every node visited in the ast tree.

visitor = FunctionVisitor(current_module=self.current_module, current_filename=self.current_filename, errors=errors)

The ast.NodeVisitor has various methods which help in only visiting the specific nodes in the ast tree. To enforce type annotations, our plugin will override, which gets called whenever a function node is visited in the ast tree. So our FunctionVisitor class will look something like this

class AnnotationErrorCodes(enum.Enum):
    ARGUMENT_ANNOTATION_MISSING = "ETA001"
    RETURN_TYPE_ANNOTATION_MISSING = "ETA002"

class FunctionVisitor(ast.NodeVisitor):
    def __init__(self, current_module: str, current_filename: str, errors: List[Tuple[int, int, str]]) -> None:
        self.current_module = current_module
        self.current_filename = current_filename
        self.errors = errors

    def is_function_arg_annotation_present(self, arg: ast.arg) -> bool:
        args_to_skip = ['self', 'cls']
        if arg.arg not in args_to_skip and arg.annotation is None:
            return False
        return True

    def is_function_return_annotation_present(self, node: ast.FunctionDef) -> bool:
        is_special_function = node.name.startswith("__") and node.name.endswith("__")
        if is_special_function is False and node.returns is None:
            return False
        return True

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
        for arg in node.args.args:
            if self.is_function_arg_annotation_present(arg) is False:
                self.errors.append((arg.lineno, arg.col_offset, f"{AnnotationErrorCodes.ARGUMENT_ANNOTATION_MISSING.value} Argument `{arg.arg}` missing type annotation in function `{node.name}`"))

        if self.is_function_return_annotation_present(node) is False:
            self.errors.append((node.lineno, node.col_offset, f"{AnnotationErrorCodes.RETURN_TYPE_ANNOTATION_MISSING.value} Function `{node.name}` missing return type annotation"))
        return node

For type annotations, we need to check the types present in the function argument and its return type.

We check if annotations are present in the arguments and the return type of a function and add it to our errors list. We’ve also defined specific error codes from missing annotations. This error code prefix ETA should be the same for all errors you raise, and each error message should start with this error code; otherwise, flake8 doesn’t recognize the error raised.

Incorrect:

"Function `add_numbers` missing return type annotation ETA002"

Correct:

"ETA002 Function `add_numbers` missing return type annotation"

That’s all you’ve successfully created a flake8 plugin 🎊 But first let’s test this out before celebrating.

Step 4: Testing our plugin.

Big brain time 🧠: We can write unit tests for our plugin and test it without even installing it 💣

We can create a string with single or multiple lines of code and convert it into ast tree using ast.parse(code) and pass it to our plugin.

Enough talk, show me how!

import ast

import pytest
from enforce_type_annotations import Plugin

class TestPlugin:
    file_name: str = "mock_filename"

    def test_when_type_annotation_are_present(self):
        code = """def test_function(mock_arg: Any) -> None:pass"""
        tree = ast.parse(code)
        plugin = Plugin(tree=tree, filename=self.file_name)
        with pytest.raises(StopIteration):
            next(plugin.run())

    def test_argument_type_annotation_not_present_in_function(self):
        code = """def test_function(mock_arg) -> None:pass"""
        tree = ast.parse(code)
        plugin = Plugin(tree=tree, filename=self.file_name)
        for lineno, col_offset, error_message, instance in plugin.run():
            assert error_message == "ETA001 Argument `mock_arg` missing type annotation in function `test_function`"

    def test_return_type_annotation_not_present_in_function(self):
        code = """def test_function(mock_arg: Any):pass
        """
        tree = ast.parse(code)
        plugin = Plugin(tree=tree, filename=self.file_name)
        for lineno, col_offset, error_message, instance in plugin.run():
            assert error_message == "ETA002 Function `test_function` missing return type annotation"

Now moving on to the next station,

Publishing the plugin on PyPI 💥

Create a setup.py file with this content

from setuptools import setup, find_packages
from enforce_type_annotations import Plugin

with open("README.md", "r") as f:
    long_description = f.read()

setup(
    name="flake8-enforce-type-annotations",
    version=Plugin.version,
    author="Your name",
    author_email="Your email",
    description="Flake8 plugin for enforcing type annotations.",
    license="MIT",
    long_description=long_description,
    long_description_content_type="text/markdown",
    keywords=["flake8", "type", "python", "enforce", "annotations"],
    url="<https://github.com/fylein/flake8-enforce-type-annotations>",
    packages=find_packages(
        include=["enforce_type_annotations.py"]
    ),
    install_requires=[
        "flake8==7.0.0"
    ],
    entry_points={
        "flake8.extension": [
            "ETA = enforce_type_annotations:Plugin",
        ],
    },
    classifiers=[
        "Topic :: Internet :: WWW/HTTP",
        "Intended Audience :: Developers",
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
    ]
)

Here, the important things are:

• name: This will be the package's name when you install it via pip.

• packages: This should include all the files that are required for your plugin to run. Keep this minimal so that the package is of smaller size when installed.

• install_requires: Any dependent library that is required to run your plugin.

• entry_points: Flake8 plugins should be added under flake8.extension and should have the format “error_code_prefix = path_to_plugin_class”. In our case, it should be “ETA = enforce_type_annotations.checker:Plugin”.

To publish packages to PyPI we’ve created a GitHub workflow that runs whenever a new release is created in GitHub and that publishes the package directly to PyPI.

Create a new workflow .github/workflows/release.yaml and get PYPI_API_TOKEN from your PyPI account.

name: Publish to PyPI.org
on:
  release:
    types: [published]
jobs:
  pypi:
    runs-on: ubuntu-latest
    steps:
      - name: Checkout
        uses: actions/checkout@v3
        with:
          fetch-depth: 0
      - run: python3 -m pip install --upgrade build && python3 -m build

      - name: Publish package
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          password: ${{ secrets.PYPI_API_TOKEN }}

How to plan and adopt a plugin incrementally in an existing code base 📈

We run flake8 checks before developers commit their changes, i.e. via a pre-commit GitHub hook. This ensures the code being pushed already conforms to standards set by flake8 rules and only runs on the files being changed by the developer. Check out pre-commit on how to add this to your project.

As we wrap this up, reaching a cleaner code is a gradual process and flake8’s extensibility to create custom rules helps in fine-tuning our code base better. Create plugins that not only identify issues but also provide insightful feedback, contributing to a codebase that is not just correct but also comprehensible. That’s it! 🎉 🥳

Until next time 👋🏼