Add proper typing

9 files changed, 64 insertions, 47 deletions

diff --git a/patternutils/commands/patterncfg.py b/patternutils/commands/patterncfg.py
index 6b2f6ec..e84a0db 100644
--- a/patternutils/commands/patterncfg.py
+++ b/patternutils/commands/patterncfg.py
@@ -15,9 +15,14 @@ else:
     has_pygments = True
 
 
-from .. import config
+from .. import config, template
 
-CURRENT_SUPPORTED_PROGRAMS = ("pmatch", "tmv", "targs", "tln")
+CURRENT_SUPPORTED_PROGRAMS: tuple[config.CommandNames, ...] = (
+    "pmatch",
+    "tmv",
+    "targs",
+    "tln",
+)
 
 
 @click.group(context_settings={"help_option_names": ("-h", "--help", "-?")})
@@ -57,7 +62,9 @@ def list() -> None:
 )
 @click.argument("command", required=True, type=click.Choice(CURRENT_SUPPORTED_PROGRAMS))
 @click.argument("name")
-def add(template_engine: Optional[str], command: str, name: str) -> None:
+def add(
+    template_engine: Optional[template.TemplateEngines], command: str, name: str
+) -> None:
 
     if command == "pmatch":
         fileext = "regex"
@@ -77,7 +84,7 @@ def add(template_engine: Optional[str], command: str, name: str) -> None:
     help="select the command from which to remove the file",
 )
 @click.argument("name")
-def remove(command: Optional[str], name: str) -> None:
+def remove(command: Optional[config.CommandNames], name: str) -> None:
 
     files = config.find_templates(name, command)
 
@@ -112,7 +119,7 @@ def remove(command: Optional[str], name: str) -> None:
     help="Colorize output using pygments",
 )
 @click.argument("name")
-def view(name: str, command: str, color: bool) -> None:
+def view(name: str, command: config.CommandNames, color: bool) -> None:
     files = config.find_templates(name, command)
 
     if not files:
@@ -162,7 +169,7 @@ def view(name: str, command: str, color: bool) -> None:
     help="select the command from which to edit the file",
 )
 @click.argument("name")
-def edit(name: str, command: str) -> None:
+def edit(name: str, command: config.CommandNames) -> None:
 
     files = config.find_templates(name, command)
     if not files:
diff --git a/patternutils/commands/pjson.py b/patternutils/commands/pjson.py
index 8604971..12469ba 100644
--- a/patternutils/commands/pjson.py
+++ b/patternutils/commands/pjson.py
@@ -1,6 +1,6 @@
 import json
 import re
-from typing import Any, Pattern
+from typing import Pattern
 
 import click
 
diff --git a/patternutils/commands/pmatch.py b/patternutils/commands/pmatch.py
index 3884a9d..44e55b4 100644
--- a/patternutils/commands/pmatch.py
+++ b/patternutils/commands/pmatch.py
@@ -1,8 +1,9 @@
 import functools
 import itertools
+import os
 import os.path
 import re
-from typing import Any, Generator, Iterator, Optional, Pattern
+from typing import Generator, Iterator, Optional, Pattern
 
 import click
 
@@ -12,7 +13,7 @@ from .. import utils
 
 
 def apply_regex(
-    walk_function: Iterator[Any],
+    walk_function: Iterator[os.DirEntry[str]],
     regex_pattern: Pattern[str],
     match_full_path: bool,
 ) -> Generator[Optional[dict[str, str]], None, None]:
diff --git a/patternutils/commands/targs.py b/patternutils/commands/targs.py
index 34cf018..a88bd71 100644
--- a/patternutils/commands/targs.py
+++ b/patternutils/commands/targs.py
@@ -77,7 +77,7 @@ def targs(
     command: str,
     read_from: click.File,
     shell: bool,
-    template_engine: str,
+    template_engine: template.TemplateEngines,
     max_procs: int,
     redirect_stdout: bool,
     redirect_stderr: bool,
diff --git a/patternutils/commands/tln.py b/patternutils/commands/tln.py
index 878f556..ece5b75 100644
--- a/patternutils/commands/tln.py
+++ b/patternutils/commands/tln.py
@@ -91,7 +91,7 @@ def tln(
     editor: bool,
     force: bool,
     symbolic: bool,
-    template_engine: str,
+    template_engine: template.TemplateEngines,
     use_predefined_template: bool,
     read_from: click.File,
     key: str,
@@ -108,9 +108,7 @@ def tln(
 
     if use_predefined_template:
         try:
-            output_pattern, template_engine = config.load_template(
-                output_pattern, "tln"
-            )
+            output_pattern, template_engine = config.load_template(output_pattern, "tln")
         except FileNotFoundError:
             click.secho(f"Template {output_pattern} not found!", fg="red", err=True)
             raise SystemExit
diff --git a/patternutils/commands/tmv.py b/patternutils/commands/tmv.py
index fa7105e..6ce45ed 100644
--- a/patternutils/commands/tmv.py
+++ b/patternutils/commands/tmv.py
@@ -80,7 +80,7 @@ def tmv(
     interactive: bool,
     editor: bool,
     force: bool,
-    template_engine: str,
+    template_engine: template.TemplateEngines,
     abort_on_path_exist: bool,
     use_predefined_template: bool,
     read_from: click.File,
@@ -106,9 +106,7 @@ def tmv(
 
     if use_predefined_template:
         try:
-            output_pattern, template_engine = config.load_template(
-                output_pattern, "tmv"
-            )
+            output_pattern, template_engine = config.load_template(output_pattern, "tmv")
         except FileNotFoundError:
             click.secho(f"Template {output_pattern} not found!", fg="red", err=True)
             raise SystemExit
diff --git a/patternutils/config.py b/patternutils/config.py
index f999b2e..b18b72c 100644
--- a/patternutils/config.py
+++ b/patternutils/config.py
@@ -1,17 +1,29 @@
 import glob
 import os.path
 import re
-from typing import Optional
+from typing import Literal, Optional, cast, overload
 
 import click
 
+TemplateExtensions = Literal["python", "jinja2", "regex"]
+CommandNames = Literal["pmatch", "tmv", "targs", "tln"]
+
 CONFIG_PATH = click.get_app_dir("patternutils")
 
 REGEX_PROGRAMS = ("pmatch",)
 TEMPLATE_PROGRAMS = ("tmv", "targs", "tln")
 
+@overload
+def load_template(template_name: str, command_name: Literal["pmatch"]) -> tuple[str, Literal["regex"]]:
+    ...
+
+@overload
+def load_template(template_name: str, command_name: Literal["tmv", "targs", "tln"]) -> tuple[str, Literal["python", "jinja2"]]:
+    ...
 
-def load_template(template_name: str, command_name: Optional[str]) -> tuple[str, str]:
+def load_template(
+    template_name: str, command_name: Optional[CommandNames]
+) -> tuple[str, TemplateExtensions]:
     """Load a template from the config directory."""
 
     files = find_templates(template_name, command_name)
@@ -25,33 +37,31 @@ def load_template(template_name: str, command_name: Optional[str]) -> tuple[str,
         raise RuntimeError("Multiple files found!")
 
     output_pattern = template_lines.replace("\n", "")
-    template_engine = os.path.splitext(template_file)[1][1:]
+    template_engine = cast(TemplateExtensions, os.path.splitext(template_file)[1][1:])
 
     return output_pattern, template_engine
 
 
-def find_templates(template_name: str, command_name: Optional[str] = None) -> List[str]:
+def find_templates(
+    template_name: str, command_name: Optional[CommandNames] = None
+) -> list[str]:
 
     if command_name is None:
         # CONFIG_PATH/*/TEMPLATE_NAME.*
         path_sub_regex = ".*"
         extension = ".*"
-    elif command_name in REGEX_PROGRAMS:
-        path_sub_regex = f"(globals|{command_name})"
-        extension = "regex"
-    elif command_name in TEMPLATE_PROGRAMS:
-        # CONFIG_PATH/globals/TEMPLATE_NAME.jinja2
-        # CONFIG_PATH/globals/TEMPLATE_NAME.python
-        # CONFIG_PATH/COMMAND_NAME/TEMPLATE_NAME.jinja2
-        # CONFIG_PATH/COMMAND_NAME/TEMPLATE_NAME.python
-        path_sub_regex = f"(globals|{command_name})"
-        extension = "(jinja2|python)"
     else:
-        # Fail over if command name is not known to have a specific extension
-        # This should normally not happen
-        # CONFIG_PATH/COMMAND_NAME/TEMPLATE_NAME.*
-        path_sub_regex = f"{command_name}"
-        extension = ".*"
+        path_sub_regex = f"(globals|{command_name})"
+        if command_name in REGEX_PROGRAMS:
+            extension = "regex"
+        elif command_name in TEMPLATE_PROGRAMS:
+            # CONFIG_PATH/globals/TEMPLATE_NAME.jinja2
+            # CONFIG_PATH/globals/TEMPLATE_NAME.python
+            # CONFIG_PATH/COMMAND_NAME/TEMPLATE_NAME.jinja2
+            # CONFIG_PATH/COMMAND_NAME/TEMPLATE_NAME.python
+            extension = "(jinja2|python)"
+        else:
+            raise ValueError(f"{command_name} is not known!")
 
     regex = re.compile(
         os.path.join(CONFIG_PATH, f"{path_sub_regex}", f"{template_name}\.{extension}")
diff --git a/patternutils/template.py b/patternutils/template.py
index f0b3c43..a05be83 100644
--- a/patternutils/template.py
+++ b/patternutils/template.py
@@ -1,4 +1,6 @@
-from typing import Callable
+from typing import Callable, Literal, cast
+
+TemplateEngines = Literal["python", "jinja2"]
 
 try:
     import jinja2
@@ -42,13 +44,14 @@ else:
     jinja_env.filters["splitext"] = splitext
 
 
-def get_render_function(template: str, *, engine: str = "python") -> Callable[..., str]:
+def get_render_function(
+    template: str, *, engine: TemplateEngines = "python"
+) -> Callable[..., str]:
     if engine == "jinja2":
         if jinja_feature:
             env_template = jinja_env.from_string(template)
-            render_function = env_template.render
+            return cast(Callable[..., str], env_template.render)
         else:
             raise RuntimeError("jinja2 not installed!")
     else:
-        render_function = template.format
-    return render_function
+        return template.format
diff --git a/patternutils/utils.py b/patternutils/utils.py
index 0395110..88cad88 100644
--- a/patternutils/utils.py
+++ b/patternutils/utils.py
@@ -1,6 +1,6 @@
 import json
 import os
-from typing import Any, Generator, Optional, Union
+from typing import Any, Generator, Optional
 
 import click
 
@@ -10,7 +10,7 @@ def max_procs() -> int:
     return 1 if cpu_count is None else cpu_count * 2
 
 
-def edit(data: Union[dict[str, str], list[str]]) -> Optional[dict[str, str]]:
+def edit(data: dict[str, str] | list[str]) -> Optional[dict[str, str]]:
 
     longest = len(max(data, key=len))
 
@@ -38,11 +38,11 @@ def edit(data: Union[dict[str, str], list[str]]) -> Optional[dict[str, str]]:
 
 
 def json_dumps(data: Any, human_readable: bool = False) -> str:
-    config = (
+    config: dict[str, Any] = (
         {"indent": 4, "sort_keys": True}
         if human_readable
         else {"separators": (",", ":")}
-    )  # type: dict[str, Any]
+    )
     return json.dumps(data, **config)
 
 
@@ -56,7 +56,7 @@ def json_loads(data: Any) -> Any:
 
 def walk(
     directory: str, recursive: bool = False, match_directories: bool = False
-) -> Generator[os.DirEntry[Any], None, None]:
+) -> Generator[os.DirEntry[str], None, None]:
     for entry in os.scandir(directory):
         is_dir = entry.is_dir()
         if recursive and is_dir: