Repository: peopledoc/septentrion

Issue:
Add type annotations
The project is Py3.6+ which means we can add type annotations everywhere. Mypy is already installed, so let's do this.

Discussion:

Patch:
diff --git a/septentrion/__init__.py b/septentrion/__init__.py
index bc30031..ba98f63 100644
--- a/septentrion/__init__.py
+++ b/septentrion/__init__.py
@@ -1,4 +1,4 @@
-def _extract_version(package_name):
+def _extract_version(package_name: str) -> str:
     try:
         # Package is installed (even with -e)
         import pkg_resources
diff --git a/septentrion/cli.py b/septentrion/cli.py
index 47c47a1..bdc033d 100644
--- a/septentrion/cli.py
+++ b/septentrion/cli.py
@@ -5,22 +5,17 @@ it.
 """
 import logging
 import os
+from typing import Any, Dict
 
 import click
 from click.types import StringParamType
 
-from septentrion import __version__
-from septentrion import core
-from septentrion import db
-from septentrion import migrate
-from septentrion import settings
-from septentrion import style
-from septentrion import utils
+from septentrion import __version__, core, db, migrate, settings, style, utils
 
 logger = logging.getLogger(__name__)
 
 
-def get_context_settings():
+def get_context_settings() -> Dict[str, Any]:
     try:
         with open("septentrion.ini") as file:
             config = file.read()
@@ -38,14 +33,15 @@ def get_context_settings():
 CONTEXT_SETTINGS = get_context_settings()
 
 
-def print_version(ctx, __, value):
+def print_version(ctx: click.Context, param: Any, value: bool):
+    # Value will be True is --version is present, false otherwise.
     if not value or ctx.resilient_parsing:
         return
     click.echo("Septentrion {}".format(__version__))
     ctx.exit()
 
 
-def validate_version(ctx, param, value):
+def validate_version(ctx: click.Context, param: Any, value: str):
     if value and not utils.is_version(value):
         raise click.BadParameter("{value} is not a valid version")
     return value
@@ -54,7 +50,7 @@ def validate_version(ctx, param, value):
 class CommaSeparatedMultipleString(StringParamType):
     envvar_list_splitter = ","
 
-    def split_envvar_value(self, rv):
+    def split_envvar_value(self, rv: str):
         values = super(CommaSeparatedMultipleString, self).split_envvar_value(rv)
         return tuple(value.strip() for value in values)
 
@@ -186,7 +182,7 @@ def migrate_func():
 
 @cli.command()
 @click.argument("version", callback=validate_version)
-def fake(version):
+def fake(version: str):
     """
     Fake migrations until version.
     Write migrations in the migration table without applying them, for
diff --git a/septentrion/core.py b/septentrion/core.py
index d3318e9..6bfd398 100644
--- a/septentrion/core.py
+++ b/septentrion/core.py
@@ -4,16 +4,13 @@ from the existing files (septentrion.files) and from the db (septentrion.db)
 """
 
 import os
+from typing import Any, Dict, Iterable, Optional
 
-from septentrion import db
-from septentrion import exceptions
-from septentrion import files
-from septentrion import style
-from septentrion import utils
+from septentrion import db, exceptions, files, style, utils
 from septentrion.settings import settings
 
 
-def get_applied_versions():
+def get_applied_versions() -> Iterable[str]:
     """
     Return the list of applied versions.
     Reuse django migration table.
@@ -25,7 +22,12 @@ def get_applied_versions():
     return utils.sort_versions(applied_versions & known_versions)
 
 
-def get_closest_version(target_version, sql_tpl, existing_files, force_version=None):
+def get_closest_version(
+    target_version: str,
+    sql_tpl: str,
+    existing_files: Iterable[str],
+    force_version: Optional[str] = None,
+):
     """
     Get the version of a file (schema or fixtures) to use to init a DB.
     Take the closest to the target_version. Can be the same version, or older.
@@ -66,7 +68,7 @@ def get_closest_version(target_version, sql_tpl, existing_files, force_version=N
     return None
 
 
-def get_best_schema_version():
+def get_best_schema_version() -> str:
     """
     Get the best candidate to init the DB.
     """
@@ -82,7 +84,7 @@ def get_best_schema_version():
     return version
 
 
-def get_fixtures_version(target_version):
+def get_fixtures_version(target_version: str) -> str:
     """
     Get the closest fixtures to use to init a new DB
     to the current target version.
@@ -98,7 +100,7 @@ def get_fixtures_version(target_version):
     return version
 
 
-def build_migration_plan():
+def build_migration_plan() -> Iterable[Dict[str, Any]]:
     """
     Return the list of migrations by version,
     from the version used to init the DB to the current target version.
@@ -133,7 +135,7 @@ def build_migration_plan():
         yield {"version": version, "plan": version_plan}
 
 
-def describe_migration_plan(stylist=style.noop_stylist):
+def describe_migration_plan(stylist: style.Stylist = style.noop_stylist) -> None:
     current_version = db.get_current_schema_version()
     with stylist.activate("title") as echo:
         echo("Current version is {}".format(current_version))
diff --git a/septentrion/db.py b/septentrion/db.py
index e96ced3..2729d10 100644
--- a/septentrion/db.py
+++ b/septentrion/db.py
@@ -5,8 +5,10 @@ Interact with the migrations table.
 import logging
 from contextlib import contextmanager
 from distutils.version import StrictVersion
+from typing import Any, Iterable, Optional, Tuple
 
 import psycopg2
+from psycopg2.extensions import connection as Connection
 from psycopg2.extras import DictCursor
 
 from septentrion.settings import settings
@@ -14,7 +16,7 @@ from septentrion.settings import settings
 logger = logging.getLogger(__name__)
 
 
-def get_connection():
+def get_connection() -> Connection:
     """
     Opens a PostgreSQL connection using psycopg2.
     """
@@ -48,7 +50,7 @@ def get_connection():
 
 
 @contextmanager
-def execute(query, args=tuple(), commit=False):
+def execute(query: str, args: Tuple = tuple(), commit: bool = False) -> Any:
     query = " ".join(query.format(table=settings.TABLE).split())
     with get_connection() as conn:
         with conn.cursor(cursor_factory=DictCursor) as cur:
@@ -60,7 +62,7 @@ def execute(query, args=tuple(), commit=False):
 
 
 class Query(object):
-    def __init__(self, query, args=tuple(), commit=False):
+    def __init__(self, query: str, args: Tuple = tuple(), commit: bool = False):
         self.context_manager = execute(query, args, commit)
 
     def __enter__(self):
@@ -99,24 +101,24 @@ query_is_schema_initialized = """
 """
 
 
-def get_current_schema_version():
+def get_current_schema_version() -> Optional[str]:
     versions = get_applied_versions()
     if not versions:
         return None
-    return max(StrictVersion(version) for version in versions)
+    return str(max(StrictVersion(version) for version in versions))
 
 
-def get_applied_versions():
+def get_applied_versions() -> Iterable[str]:
     with Query(query_max_version) as cur:
         return [row[0] for row in cur]
 
 
-def get_applied_migrations(version):
+def get_applied_migrations(version: str) -> Iterable[str]:
     with Query(query_get_applied_migrations, (version,)) as cur:
         return [row[0] for row in cur]
 
 
-def is_schema_initialized():
+def is_schema_initialized() -> bool:
     with Query(query_is_schema_initialized) as cur:
         try:
             return next(cur)
@@ -124,9 +126,9 @@ def is_schema_initialized():
             return False
 
 
-def create_table():
+def create_table() -> None:
     Query(query_create_table, commit=True)()
 
 
-def write_migration(version, name):
+def write_migration(version: str, name: str) -> None:
     Query(query_write_migration, (version, name), commit=True)()
diff --git a/septentrion/files.py b/septentrion/files.py
index 8f8a0a0..751380d 100644
--- a/septentrion/files.py
+++ b/septentrion/files.py
@@ -5,20 +5,21 @@ Interact with the migration files.
 import io
 import os
 from distutils.version import StrictVersion
+from typing import Iterable
 
 from septentrion import utils
 from septentrion.settings import settings
 
 
-def list_dirs(root):
+def list_dirs(root: str) -> Iterable[str]:
     return [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
 
 
-def list_files(root):
+def list_files(root: str) -> Iterable[str]:
     return [d for d in os.listdir(root) if os.path.isfile(os.path.join(root, d))]
 
 
-def get_known_versions():
+def get_known_versions() -> Iterable[str]:
     """
     Return the list of the known versions defined in migration repository,
     ordered.
@@ -43,7 +44,7 @@ def get_known_versions():
     return versions
 
 
-def is_manual_migration(migration_path):
+def is_manual_migration(migration_path: str) -> bool:
 
     if "/manual/" in migration_path:
         return True
@@ -59,18 +60,18 @@ def is_manual_migration(migration_path):
     return False
 
 
-def get_known_schemas():
+def get_known_schemas() -> Iterable[str]:
     return os.listdir(os.path.join(settings.MIGRATIONS_ROOT, "schemas"))
 
 
-def get_known_fixtures():
+def get_known_fixtures() -> Iterable[str]:
     try:
         return os.listdir(os.path.join(settings.MIGRATIONS_ROOT, "fixtures"))
     except FileNotFoundError:
         return []
 
 
-def get_migrations_files_mapping(version):
+def get_migrations_files_mapping(version: str):
     """
     Return an dict containing the list of migrations for
     the given version.
@@ -78,7 +79,7 @@ def get_migrations_files_mapping(version):
     Value: path to the migration file.
     """
 
-    def filter_migrations(files):
+    def filter_migrations(files: Iterable[str]) -> Iterable[str]:
         return [f for f in files if f.endswith("ddl.sql") or f.endswith("dml.sql")]
 
     version_root = os.path.join(settings.MIGRATIONS_ROOT, version)
diff --git a/septentrion/migrate.py b/septentrion/migrate.py
index 69f2bbe..861bbfd 100644
--- a/septentrion/migrate.py
+++ b/septentrion/migrate.py
@@ -3,19 +3,13 @@ import io
 import logging
 import os.path
 
-from septentrion import core
-from septentrion import db
-from septentrion import exceptions
-from septentrion import files
-from septentrion import runner
-from septentrion import style
-from septentrion import utils
+from septentrion import core, db, exceptions, files, runner, style, utils
 from septentrion.settings import settings
 
 logger = logging.getLogger(__name__)
 
 
-def migrate(stylist=style.noop_stylist):
+def migrate(stylist: style.Stylist = style.noop_stylist) -> None:
 
     logger.info("Starting migrations")
     if not db.is_schema_initialized():
@@ -56,7 +50,7 @@ def migrate(stylist=style.noop_stylist):
                     db.write_migration(version, mig)
 
 
-def init_schema(init_version, stylist=style.noop_stylist):
+def init_schema(init_version: str, stylist: style.Stylist = style.noop_stylist) -> None:
     # load additional files
     logger.info("Looking for additional files to run")
     additional_files = settings.ADDITIONAL_SCHEMA_FILE
@@ -107,7 +101,9 @@ def init_schema(init_version, stylist=style.noop_stylist):
         logger.info("Not applying fixtures: %s", exception)
 
 
-def create_fake_entries(version, stylist=style.noop_stylist):
+def create_fake_entries(
+    version: str, stylist: style.Stylist = style.noop_stylist
+) -> None:
     """
     Write entries in the migration table for all existing migrations
     up until the given version (included).
@@ -134,8 +130,8 @@ def create_fake_entries(version, stylist=style.noop_stylist):
                 db.write_migration(version, migration_name)
 
 
-def run_script(path):
+def run_script(path: str) -> None:
     logger.info("Running SQL file %s", path)
     with io.open(path, "r", encoding="utf8") as f:
-        script = runner.Script(f)
+        script = runner.Script(f, f.name)
         script.run(db.get_connection())
diff --git a/septentrion/runner.py b/septentrion/runner.py
index a2f6cfe..434eab4 100644
--- a/septentrion/runner.py
+++ b/septentrion/runner.py
@@ -1,6 +1,8 @@
 import logging
+from typing import Iterable
 
 import sqlparse
+from psycopg2.extensions import cursor as Cursor
 
 from septentrion import files
 from septentrion.settings import settings
@@ -12,7 +14,7 @@ class SQLRunnerException(Exception):
     pass
 
 
-def clean_sql_code(code):
+def clean_sql_code(code: str) -> str:
     output = ""
     for line in code.split("\n"):
         stripped_line = line.strip()
@@ -29,17 +31,17 @@ class Block(object):
         self.closed = False
         self.content = ""
 
-    def append_line(self, line):
+    def append_line(self, line: str) -> None:
         if self.closed:
             raise SQLRunnerException("Block closed !")
         self.content += line
 
-    def close(self):
+    def close(self) -> None:
         if self.closed:
             raise SQLRunnerException("Block closed !")
         self.closed = True
 
-    def run(self, cursor):
+    def run(self, cursor: Cursor) -> int:
         statements = sqlparse.parse(self.content)
 
         text_type = type(u"")  # Remove when only PY3 is supported
@@ -67,13 +69,13 @@ class SimpleBlock(Block):
 
 
 class MetaBlock(Block):
-    def __init__(self, command):
+    def __init__(self, command: str):
         super(MetaBlock, self).__init__()
         self.command = command
         if command != "do-until-0":
             raise SQLRunnerException("Unexpected command {}".format(command))
 
-    def run(self, cursor):
+    def run(self, cursor: Cursor) -> int:
         total_rows = 0
         # Simply call super().run in a loop...
         delta = 0
@@ -90,8 +92,8 @@ class MetaBlock(Block):
 
 
 class Script(object):
-    def __init__(self, file_handler):
-        is_manual = files.is_manual_migration(file_handler.name)
+    def __init__(self, file_handler: Iterable[str], name: str):
+        is_manual = files.is_manual_migration(name)
         if is_manual:
             self.block_list = [Block()]
         elif self.contains_non_transactional_keyword(file_handler):
diff --git a/septentrion/settings.py b/septentrion/settings.py
index f0b8818..f4f9b03 100644
--- a/septentrion/settings.py
+++ b/septentrion/settings.py
@@ -4,11 +4,12 @@ the settings, see cli.py (for now)
 """
 import configparser
 import logging
+from typing import Any, Dict
 
 logger = logging.getLogger(__name__)
 
 
-def get_config_settings(content):
+def get_config_settings(content: str) -> Dict:
     """
     Read configuration file and return
     a dict with values to use if they are
@@ -18,7 +19,7 @@ def get_config_settings(content):
     parser.read_string(content)
 
     if "septentrion" in parser:
-        config = dict(parser["septentrion"])
+        config: Dict[str, Any] = dict(parser["septentrion"])
         if "additional_schema_file" in config:
             config["additional_schema_file"] = [
                 line for line in config["additional_schema_file"].splitlines() if line
@@ -32,21 +33,22 @@ def get_config_settings(content):
     return {}
 
 
-class Settings(object):
+class Settings(dict):
+    def __getattr__(self, key: str) -> Any:
+        return self[key]
+
     @staticmethod
-    def _clean_key(key):
+    def _clean_key(key: str) -> str:
         return key.upper()
 
-    def set(self, **kwargs):
+    def set(self, **kwargs) -> None:
         """
         TODO: clean this.
         """
-        kwargs = {self._clean_key(key): value for key, value in kwargs.items()}
-
-        vars(self).update(kwargs)
+        self.update({self._clean_key(key): value for key, value in kwargs.items()})
 
 
-def log_level(verbosity):
+def log_level(verbosity: int) -> int:
     """
     verbosity  | log level
     (input)    | (output)
@@ -59,7 +61,7 @@ def log_level(verbosity):
     return 40 - 10 * min(verbosity, 3)
 
 
-def consolidate(**kwargs):
+def consolidate(**kwargs) -> None:
     settings.set(**kwargs)
 
     level = log_level(settings.VERBOSE)
diff --git a/septentrion/style.py b/septentrion/style.py
index de48c8a..ac730ba 100644
--- a/septentrion/style.py
+++ b/septentrion/style.py
@@ -1,6 +1,6 @@
 from collections import defaultdict
 from contextlib import contextmanager
-from typing import Dict
+from typing import Any, Dict, Optional
 
 import colorama
 
@@ -8,7 +8,7 @@ colorama.init()
 
 
 class Stylist(object):
-    styles = {
+    styles: Dict[str, str] = {
         "reset": colorama.Style.RESET_ALL,
         "title": colorama.Fore.CYAN + colorama.Style.BRIGHT,
         "subtitle": colorama.Fore.CYAN,
@@ -16,25 +16,27 @@ class Stylist(object):
         "danger": colorama.Fore.RED,
     }
 
-    def echo(self, content="", nl=True):
+    def echo(self, content: str = "", nl: bool = True) -> None:
         print(content, end="\n" if nl else "", flush=True)
 
     @contextmanager
-    def activate(self, style):
+    def activate(self, style: str) -> Any:
         self.echo(self.styles[style], nl=False)
         try:
             yield self.echo
         finally:
             self.echo(self.styles["reset"], nl=False)
 
-    def draw_checkbox(self, content, checked, margin=2):
+    def draw_checkbox(self, content: str, checked: bool, margin: int = 2) -> None:
         self.echo(" " * margin, nl=False)
         with self.activate("success" if checked else "danger") as echo:
             echo("[{}] ".format("X" if checked else " "), nl=False)
         echo(content, nl=False)
 
     @contextmanager
-    def checkbox(self, content, content_after=None, margin=2):
+    def checkbox(
+        self, content: str, content_after: Optional[str] = None, margin: int = 2
+    ) -> Any:
         content_after = content_after or content
         space_len = max(0, len(content) - len(content_after))
         self.draw_checkbox(content=content, checked=False, margin=margin)
diff --git a/septentrion/utils.py b/septentrion/utils.py
index fcc8cd7..03d2073 100644
--- a/septentrion/utils.py
+++ b/septentrion/utils.py
@@ -3,9 +3,10 @@ All functions in here should be easily unit testable
 """
 
 from distutils.version import StrictVersion
+from typing import Iterable, TypeVar
 
 
-def sort_versions(iterable):
+def sort_versions(iterable: Iterable[str]) -> Iterable[str]:
     """
     Sorts an iterable of strings by increasing
     equivalent version number
@@ -15,7 +16,7 @@ def sort_versions(iterable):
     return sorted(iterable, key=StrictVersion)
 
 
-def is_version(vstring):
+def is_version(vstring: str) -> bool:
     """
     Returns True if vstring is a valid version descriptor
     >>> is_version("1.2")
@@ -30,7 +31,10 @@ def is_version(vstring):
     return True
 
 
-def until(iterable, value):
+T = TypeVar("T")
+
+
+def until(iterable: Iterable[T], value: T) -> Iterable[T]:
     """
     Returns the values from iterable up until element is found
     >>> list(until(range(300), 3))
diff --git a/setup.cfg b/setup.cfg
index c0c5775..ba3b520 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -80,7 +80,6 @@ force_grid_wrap=0
 combine_as_imports=True
 line_length=88
 known_first_party=septentrion
-force_single_line=True
 not_skip = __init__.py
 
 [tool:pytest]


Test patch:
diff --git a/tests/integration/test_cli.py b/tests/integration/test_cli.py
index 7f87797..8bfd250 100644
--- a/tests/integration/test_cli.py
+++ b/tests/integration/test_cli.py
@@ -1,6 +1,5 @@
 from septentrion.__main__ import main
-from septentrion.db import get_current_schema_version
-from septentrion.db import is_schema_initialized
+from septentrion.db import get_current_schema_version, is_schema_initialized
 
 
 def test_version(cli_runner):
diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py
index e8a01e6..fd1aef8 100644
--- a/tests/unit/test_core.py
+++ b/tests/unit/test_core.py
@@ -1,8 +1,6 @@
 import pytest
 
-from septentrion import core
-from septentrion import exceptions
-from septentrion import settings
+from septentrion import core, exceptions, settings
 
 
 @pytest.fixture
diff --git a/tests/unit/test_files.py b/tests/unit/test_files.py
index dc2c4f5..8c9af46 100644
--- a/tests/unit/test_files.py
+++ b/tests/unit/test_files.py
@@ -2,8 +2,7 @@ import io
 
 import pytest
 
-from septentrion import files
-from septentrion import settings
+from septentrion import files, settings
 
 
 @pytest.mark.parametrize("isdir,expected", [(True, ["15.0"]), (False, [])])
diff --git a/tests/unit/test_runner.py b/tests/unit/test_runner.py
index a90117f..1e1e526 100644
--- a/tests/unit/test_runner.py
+++ b/tests/unit/test_runner.py
@@ -3,8 +3,7 @@ import logging
 
 import pytest
 
-from septentrion import runner
-from septentrion import settings
+from septentrion import runner, settings
 
 
 def test_clean_sql_code():
@@ -122,11 +121,10 @@ from author_data;
 select true;
 """
     )
-    handler.name = "/manual/foo.sql"
     settings.consolidate(
         non_transactional_keyword=["CONCURRENTLY", "ALTER TYPE", "VACUUM"], verbose=0
     )
-    script = runner.Script(handler)
+    script = runner.Script(handler, "/manual/foo.sql")
 
     b1, b2, b3 = script.block_list
 
diff --git a/tests/unit/test_style.py b/tests/unit/test_style.py
index db23bb0..9776fa9 100644
--- a/tests/unit/test_style.py
+++ b/tests/unit/test_style.py
@@ -1,7 +1,6 @@
 import pytest
 
-from septentrion.style import noop_stylist
-from septentrion.style import stylist
+from septentrion.style import noop_stylist, stylist
 
 CHECKED = "\x1b[32m[X] \x1b[0m"
 NOT_CHECKED = "\x1b[31m[ ] \x1b[0m"
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index da9fdc7..52dc0f4 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -1,8 +1,6 @@
 import pytest
 
-from septentrion.utils import is_version
-from septentrion.utils import sort_versions
-from septentrion.utils import until
+from septentrion.utils import is_version, sort_versions, until
 
 
 def test_sort_versions():