A flake8 plugin to enforce specific import policies throughout Python codebase.
from __future__ import annotations

import typing


class SourceConfig(typing.NamedTuple):
    allow_absolute: bool = False
    allow_from_module: bool = False
    allow_from_member: bool = False

    def asdict(self) -> dict[str, bool]:
        return self._asdict()


class Override(typing.NamedTuple):
    allow_absolute: bool | None = None
    allow_from_module: bool | None = None
    allow_from_member: bool | None = None

    def omit_nones(self) -> dict[str, bool]:
        return {k: v for k, v in self._asdict().items() if v is not None}

    def evolve(self, **update: bool) -> Override:
        return Override(**{**self.omit_nones(), **update})

    def __or__(self, source_config: SourceConfig) -> SourceConfig:
        return SourceConfig(**{**source_config.asdict(), **self.omit_nones()})


class Config(typing.NamedTuple):
    overrides: typing.Mapping[str, Override]
    registered_aliases: typing.Mapping[str, str]
    check_init: bool = False

    future: SourceConfig = SourceConfig(
        allow_absolute=True, allow_from_module=True, allow_from_member=True
    )
    stdlib: SourceConfig = SourceConfig(allow_absolute=True)
    third_party: SourceConfig = SourceConfig(allow_absolute=True)
    first_party: SourceConfig = SourceConfig(
        allow_absolute=True, allow_from_module=True
    )

    max_relative_level: int = 1
    allow_relative_from_module: bool = True
    allow_relative_from_member: bool = False

flake8_import_policy/config.py
albshady/flake8-import-policy
from __future__ import annotations

import ast
import collections
import enum
import importlib.metadata
import importlib.util
import pathlib
import sys
import typing

import flake8.options.manager  # type: ignore[import]
import isort
import isort.sections

from . import config


FUTURE_IMPORT_VIOLATION = "FIP000 __future__ module import policy violation"
STDLIB_IMPORT_VIOLATION = "FIP001 `{}` is forbidden"
THIRD_PARTY_IMPORT_VIOLATION = "FIP002 `{}` is forbidden"
FIRST_PARTY_IMPORT_VIOLATION = "FIP003 `{}` is forbidden"
RELATIVE_IMPORT_VIOLATION = "FIP004 `{}` is forbidden"
UNREGISTERED_ALIAS_ABUSE = (
    "FIP005 use of unregistered alias: `{module_name}` -> `{alias}`"
)


class SourceType(enum.Enum):
    FUTURE = enum.auto()
    STDLIB = enum.auto()
    THIRD_PARTY = enum.auto()
    FIRST_PARTY = enum.auto()


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

    _config: config.Config

    def __init__(self, tree: ast.AST, filename: str) -> None:
        self._tree = tree
        self._filepath = pathlib.Path(filename)
        self._config_by_source = {
            SourceType.FUTURE: self._config.future,
            SourceType.STDLIB: self._config.stdlib,
            SourceType.THIRD_PARTY: self._config.third_party,
            SourceType.FIRST_PARTY: self._config.first_party,
        }
        self._error_by_source = {
            SourceType.FUTURE: FUTURE_IMPORT_VIOLATION,
            SourceType.STDLIB: STDLIB_IMPORT_VIOLATION,
            SourceType.THIRD_PARTY: THIRD_PARTY_IMPORT_VIOLATION,
            SourceType.FIRST_PARTY: FIRST_PARTY_IMPORT_VIOLATION,
        }

    @classmethod
    def add_options(cls, parser: flake8.options.manager.OptionManager) -> None:
        parser.add_option(
            '--forbid-stdlib-absolute',
            action='store_true',
            default=False,
            parse_from_config=True,
            help="Forbid absolute imports for stdlib",
        )
        parser.add_option(
            '--allow-stdlib-from-module',
            action='store_true',
            default=False,
            parse_from_config=True,
            help="Allow `from module import ...` for stdlib",
        )

        parser.add_option(
            '--forbid-third-party-absolute',
            action='store_true',
            default=False,
            parse_from_config=True,
            help="Forbid absolute imports for third-party",
        )
        parser.add_option(
            '--allow-third-party-from-module',
            action='store_true',
            default=False,
            parse_from_config=True,
            help="Allow `from module import ...` for third-party",
        )

        parser.add_option(
            '--forbid-local-absolute',
            action='store_true',
            default=False,
            parse_from_config=True,
            help="Forbid absolute imports for local modules",
        )
        parser.add_option(
            '--forbid-local-from-module',
            action='store_true',
            default=False,
            parse_from_config=True,
            help="Forbid `from module import ...` for local modules",
        )

        parser.add_option(
            '--max-relative-level',
            type='int',
            default=1,
            parse_from_config=True,
            help="Max allowed level for relative imoprts (e.g. 1 for `.`, 2 for `..`, etc.)",
        )
        parser.add_option(
            '--forbid-relative-from-module',
            action='store_true',
            default=False,
            parse_from_config=True,
            help="Forbid `from .module import ...` for relative modules",
        )
        parser.add_option(
            '--register-import-aliases',
            default='',
            dest='registered_import_aliases',
            parse_from_config=True,
            comma_separated_list=True,
            help=(
                "Register allowed aliases for modules. "
                "Format `original=alias`, e.g. `sqlalchemy=sa,matplotlib.pyplot=plt`"
            ),
        )
        parser.add_option(
            '--allow-from-module',
            default='',
            dest='allow_from_module',
            parse_from_config=True,
            comma_separated_list=True,
            help="List of modules to always allow `from module import ...`",
        )
        parser.add_option(
            '--forbid-from-module',
            default='',
            dest='forbid_from_module',
            parse_from_config=True,
            comma_separated_list=True,
            help="List of modules to always forbid `from module import ...`",
        )
        parser.add_option(
            '--allow-absolute',
            default='',
            dest='allow_absolute',
            parse_from_config=True,
            comma_separated_list=True,
            help="List of modules to always allow `import module.something`",
        )
        parser.add_option(
            '--forbid-absolute',
            default='',
            dest='forbid_absolute',
            parse_from_config=True,
            comma_separated_list=True,
            help="List of modules to always forbid `import module.something`",
        )
        parser.add_option(
            '--init-must-follow-import-policy',
            action='store_true',
            default=False,
            parse_from_config=True,
            help="Whether `__init__.py` shall follow import policies",
        )

    @classmethod
    def parse_options(cls, options: flake8.options.manager.OptionManager) -> None:
        allow_from_module = set(options.allow_from_module)
        forbid_from_module = set(options.forbid_from_module)
        if allowed_and_forbidden := allow_from_module & forbid_from_module:
            raise ValueError(
                f"Can't simultaniously allow and forbid "
                f"from_module import for: {allowed_and_forbidden}"
            )
        allow_absolute = set(options.allow_absolute)
        forbid_absolute = set(options.forbid_absolute)
        if allowed_and_forbidden := allow_absolute & forbid_absolute:
            raise ValueError(
                f"Can't simultaniously allow and forbid "
                f"absolute import for: {allowed_and_forbidden}"
            )
        overrides: collections.defaultdict[
            str, config.Override
        ] = collections.defaultdict(config.Override)
        for module in allow_from_module:
            overrides[module] = overrides[module].evolve(allow_from_module=True)
        for module in forbid_from_module:
            overrides[module] = overrides[module].evolve(allow_from_module=False)
        for module in allow_absolute:
            overrides[module] = overrides[module].evolve(allow_absolute=True)
        for module in forbid_absolute:
            overrides[module] = overrides[module].evolve(allow_absolute=False)

        registered_aliases: dict[str, str] = {}
        for raw_alias in options.registered_import_aliases:
            full_module_path, _, alias = raw_alias.partition('=')
            registered_aliases[full_module_path] = alias

        cls._config = config.Config(
            overrides=overrides,
            registered_aliases=registered_aliases,
            check_init=options.init_must_follow_import_policy,
            stdlib=config.SourceConfig(
                allow_absolute=not options.forbid_stdlib_absolute,
                allow_from_module=options.allow_stdlib_from_module,
            ),
            third_party=config.SourceConfig(
                allow_absolute=not options.forbid_third_party_absolute,
                allow_from_module=options.allow_third_party_from_module,
            ),
            first_party=config.SourceConfig(
                allow_absolute=not options.forbid_local_absolute,
                allow_from_module=not options.forbid_local_from_module,
            ),
            max_relative_level=options.max_relative_level,
            allow_relative_from_module=not options.forbid_relative_from_module,
        )

    def run(self) -> typing.Iterator[tuple[int, int, str, type[Plugin]]]:
        sys.path.append(str(pathlib.Path.cwd()))
        if self._filepath.name.startswith('.'):
            return
        if self._filepath.name == '__init__.py' and not self._config.check_init:
            return
        for node in ast.walk(self._tree):
            if isinstance(node, ast.Import):
                errors = self._check_absolute_import(node)
            elif isinstance(node, ast.ImportFrom):
                errors = self._check_import_from(node)
            else:
                continue
            for msg in errors:
                yield (node.lineno, node.col_offset, msg, type(self))

    def _check_absolute_import(self, node: ast.Import) -> typing.Iterator[str]:
        for imported_module in node.names:
            module_name = imported_module.name
            alias = imported_module.asname
            source_config, error_template = self._get_config_and_error_template(
                module_name
            )
            if not source_config.allow_absolute:
                yield error_template.format(module_name)
            if alias is not None:
                yield from self._check_alias(module_name, alias)

    def _check_alias(self, module_name: str, alias: str) -> typing.Iterator[str]:
        if self._config.registered_aliases.get(module_name) != alias:
            yield UNREGISTERED_ALIAS_ABUSE.format(module_name=module_name, alias=alias)

    def _determine_source_type(self, module_name: str) -> SourceType:
        isort_section = isort.place_module(module_name)
        return {
            isort.sections.FUTURE: SourceType.FUTURE,
            isort.sections.STDLIB: SourceType.STDLIB,
            isort.sections.THIRDPARTY: SourceType.THIRD_PARTY,
            isort.sections.FIRSTPARTY: SourceType.FIRST_PARTY,
            isort.sections.LOCALFOLDER: SourceType.FIRST_PARTY,
        }[isort_section]

    def _check_import_from(self, node: ast.ImportFrom) -> typing.Iterator[str]:
        if node.level > 0:
            yield from self._check_relative_import(node)
            return
        assert (
            node.module is not None
        ), "node.module can't be None when import is not relative"
        source_module = node.module
        source_config, error_template = self._get_config_and_error_template(
            source_module=source_module
        )
        for imported_object_alias in node.names:
            imported_object = imported_object_alias.name
            if imported_object == '*':
                # Let flake8 handle wildcard import itself
                return
            full_imported_object_path = f'{source_module}.{imported_object}'
            if not source_config.allow_from_module:
                yield error_template.format(
                    f"from {source_module} import {imported_object}"
                )
            if (alias := imported_object_alias.asname) is not None:
                yield from self._check_alias(
                    module_name=full_imported_object_path, alias=alias
                )

    def _get_config_and_error_template(
        self, source_module: str
    ) -> tuple[config.SourceConfig, str]:
        source_type = self._determine_source_type(module_name=source_module)
        source_config = self._config_by_source[source_type]
        override = self._config.overrides.get(source_module, config.Override())
        source_config = override | source_config
        return source_config, self._error_by_source[source_type]

    def _check_relative_import(self, node: ast.ImportFrom) -> typing.Iterator[str]:
        assert node.level > 0, "Relative import's node.level must not be 0"
        if node.level > self._config.max_relative_level:
            return iter(RELATIVE_IMPORT_VIOLATION)
        source_module = str(self._filepath.parents[node.level - 1]).replace('/', '.')
        if node.module is not None:
            source_module = f'{source_module}.{node.module}'

        for imported_object_alias in node.names:
            imported_object = imported_object_alias.name
            if imported_object == '*':
                # Let flake8 handle wildcard import itself
                return
            if not self._config.allow_relative_from_module:
                yield RELATIVE_IMPORT_VIOLATION.format(
                    f"from {'.' * node.level}{node.module or ''} import {imported_object}"
                )
            full_imported_object_path = f'{source_module}.{imported_object}'
            if (alias := imported_object_alias.asname) is not None:
                yield from self._check_alias(
                    module_name=full_imported_object_path, alias=alias
                )

flake8_import_policy/plugin.py
albshady/flake8-import-policy
import pathlib

import pytest


@pytest.fixture
def data_dir() -> pathlib.Path:
    tests_dir = pathlib.Path(__file__).parent
    return tests_dir / 'data'

tests/conftest.py
albshady/flake8-import-policy
from .module import member


__all__ = ['member']

tests/local_package/__init__.py
albshady/flake8-import-policy
member = object()

tests/local_package/module.py
albshady/flake8-import-policy
from __future__ import annotations

import ast
import textwrap

import pytest

import flake8_import_policy
from flake8_import_policy import config


def create_config(
    registered_aliases: dict[str, str] | None = None,
    overrides: dict[str, config.Override] | None = None,
) -> config.Config:
    return config.Config(
        registered_aliases=registered_aliases or {},
        overrides=overrides or {},
    )


def get_errors(
    code: str,
    filename: str = 'tests/test_import_policy.py',
    plugin_config: config.Config | None = None,
) -> set[str]:
    tree = ast.parse(code)
    flake8_import_policy.Plugin._config = plugin_config or create_config()
    plugin = flake8_import_policy.Plugin(tree=tree, filename=filename)
    return {f'{line}:{col} {msg}' for line, col, msg, _ in plugin.run()}


def test_correct_stdlib_import():
    code = textwrap.dedent(
        """\
        import asyncio
        import time
        import typing
        """
    )
    errors = get_errors(code)
    assert not errors


def test_override_stdlib_import():
    code = textwrap.dedent(
        """\
        import datetime
        from datetime import datetime
        """
    )
    errors = get_errors(
        code,
        plugin_config=create_config(
            overrides={'datetime': config.Override(allow_from_module=True)}
        ),
    )
    assert not errors


def test_stdlib_import_policy_violation():
    code = "from os import path"
    errors = get_errors(code)
    assert errors == {"1:0 FIP001 `from os import path` is forbidden"}


def test_correct_third_party_import():
    code = textwrap.dedent(
        """\
        import isort.main
        import pytest
        """
    )
    errors = get_errors(code)
    assert not errors


def test_third_party_import_policy_violation():
    code = "from pytest import fixture"
    errors = get_errors(code)
    assert errors == {"1:0 FIP002 `from pytest import fixture` is forbidden"}


def test_correct_local_module_import():
    code = textwrap.dedent(
        """\
        import flake8_import_policy
        import flake8_import_policy.plugin
        from flake8_import_policy import plugin
        from flake8_import_policy import config, plugin
        """
    )
    errors = get_errors(code)
    assert not errors


@pytest.mark.xfail(reason="from_member is not supported")
def test_forbid_local_module_member_import():
    code = textwrap.dedent(
        """\
        from flake8_import_policy import Plugin
        """
    )
    errors = get_errors(code)
    assert errors == {
        "1:0 FIP003 `from flake8_import_policy import Plugin` is forbidden"
    }


@pytest.mark.xfail(reason="from_member is not supported")
def test_forbid_import_member_when_importing_multiple_from_local_module():
    code = textwrap.dedent(
        """\
        from flake8_import_policy import plugin, Plugin
        """
    )
    errors = get_errors(code)
    assert errors == {
        "1:0 FIP003 `from flake8_import_policy import Plugin` is forbidden"
    }


def test_relative_module_import():
    code = textwrap.dedent(
        """\
        from .local_package import module
        """
    )
    errors = get_errors(code)
    assert not errors


@pytest.mark.xfail(reason="from_member is not supported")
def test_forbid_import_member_from_relative_module():
    code = textwrap.dedent(
        """\
        from .local_package import member
        """
    )
    errors = get_errors(code)
    assert errors == {"1:0 FIP004 `from .local_package import member` is forbidden"}


def test_forbid_absolute_alias():
    code = textwrap.dedent(
        """\
        import datetime as dt
        """
    )
    errors = get_errors(code)
    assert errors == {"1:0 FIP005 use of unregistered alias: `datetime` -> `dt`"}


def test_allow_absolute_alias():
    code = textwrap.dedent(
        """\
        import datetime as dt
        """
    )
    errors = get_errors(
        code, plugin_config=create_config(registered_aliases={'datetime': 'dt'})
    )
    assert not errors


def test_allow_alias_but_forbid_from_member():
    code = textwrap.dedent(
        """\
        from datetime import datetime as dt
        """
    )
    errors = get_errors(
        code,
        plugin_config=create_config(registered_aliases={'datetime.datetime': 'dt'}),
    )
    assert errors == {'1:0 FIP001 `from datetime import datetime` is forbidden'}


def test_allow_from_member_but_forbid_alias():
    code = textwrap.dedent(
        """\
        from datetime import datetime as dt
        """
    )
    errors = get_errors(
        code,
        plugin_config=create_config(
            overrides={'datetime': config.Override(allow_from_module=True)},
        ),
    )
    assert errors == {
        '1:0 FIP005 use of unregistered alias: `datetime.datetime` -> `dt`'
    }


def test_allow_from_member_alias():
    code = textwrap.dedent(
        """\
        from datetime import datetime as dt
        """
    )
    errors = get_errors(
        code,
        plugin_config=create_config(
            registered_aliases={'datetime.datetime': 'dt'},
            overrides={'datetime': config.Override(allow_from_module=True)},
        ),
    )
    assert not errors


def test_forbid_relative_alias():
    code = textwrap.dedent(
        """\
        from .local_package import module as alias
        """
    )
    errors = get_errors(code)
    assert errors == {
        '1:0 FIP005 use of unregistered alias: `tests.local_package.module` -> `alias`'
    }


def test_allow_relative_alias():
    code = textwrap.dedent(
        """\
        from .local_package import module as alias
        """
    )
    errors = get_errors(
        code,
        plugin_config=create_config(
            registered_aliases={'tests.local_package.module': 'alias'}
        ),
    )
    assert not errors

tests/test_import_policy.py
albshady/flake8-import-policy