Make typing more clear and use an overlay environment for jinja2 to load files from the config folder

7 files changed, 86 insertions, 37 deletions

diff --git a/patternutils/commands/patterncfg.py b/patternutils/commands/patterncfg.py
index 3c75711..ee71d86 100644
--- a/patternutils/commands/patterncfg.py
+++ b/patternutils/commands/patterncfg.py
@@ -1,7 +1,7 @@
 import itertools
 import os
 import os.path
-from typing import Literal, Optional
+from typing import Literal, Optional, get_args
 
 import click
 
@@ -15,13 +15,16 @@ else:
     has_pygments = True
 
 
-from .. import config, template
+from .. import config
+from ..types import (
+    PredefinedCommands,
+    RegexCommands,
+    TemplateCommands,
+    TemplateExtensions,
+)
 
-CURRENT_SUPPORTED_PROGRAMS: tuple[config.CommandNames, ...] = (
-    "pmatch",
-    "tmv",
-    "targs",
-    "tln",
+CURRENT_SUPPORTED_PROGRAMS: tuple[RegexCommands | TemplateCommands, ...] = tuple(
+    itertools.chain(get_args(RegexCommands), get_args(TemplateCommands))
 )
 
 
@@ -46,7 +49,7 @@ def list() -> None:
             click.secho(directory, fg="blue")
             for filename in os.listdir(os.path.join(config.CONFIG_PATH, directory)):
                 click.echo(f"└── {filename}")
-    except FileNotFoundError as e:
+    except FileNotFoundError:
         click.secho(
             f"Config directory '{config.CONFIG_PATH}' not found! Run 'patterncfg init' first.",
             fg="red",
@@ -66,13 +69,13 @@ def list() -> None:
 @click.argument("command", required=True, type=click.Choice(CURRENT_SUPPORTED_PROGRAMS))
 @click.argument("name")
 def add(
-    template_engine: Optional[template.TemplateEngines],
+    template_engine: Optional[TemplateExtensions],
     command: str,
     name: str,
 ) -> None:
     """Creates a new template file for a given command."""
 
-    fileext: template.TemplateEngines | Literal["regex"] | None = template_engine
+    fileext: TemplateExtensions | Literal["regex"] | None = template_engine
     if command == "pmatch":
         fileext = "regex"
 
@@ -89,7 +92,7 @@ def add(
     help="Specify the command from which to remove the file.",
 )
 @click.argument("name")
-def remove(command: Optional[config.CommandNames], name: str) -> None:
+def remove(command: Optional[PredefinedCommands], name: str) -> None:
     """Removes a specified template file."""
     files = config.find_templates(name, command)
 
@@ -128,7 +131,7 @@ def remove(command: Optional[config.CommandNames], name: str) -> None:
     help="Colorize output using pygments if available.",
 )
 @click.argument("name")
-def view(name: str, command: config.CommandNames, color: bool) -> None:
+def view(name: str, command: PredefinedCommands, color: bool) -> None:
     """Displays the contents of a template file."""
 
     files = config.find_templates(name, command)
@@ -182,7 +185,7 @@ def view(name: str, command: config.CommandNames, color: bool) -> None:
     help="Specify the command from which to edit the file.",
 )
 @click.argument("name")
-def edit(name: str, command: config.CommandNames) -> None:
+def edit(name: str, command: PredefinedCommands) -> None:
     """Opens a template file in an editor."""
 
     files = config.find_templates(name, command)
diff --git a/patternutils/commands/targs.py b/patternutils/commands/targs.py
index 85acade..8280faa 100644
--- a/patternutils/commands/targs.py
+++ b/patternutils/commands/targs.py
@@ -8,6 +8,7 @@ from typing import Any, Callable, TextIO
 import click
 
 from .. import config, template, utils
+from ..types import TemplateExtensions
 
 
 def run_command(
@@ -105,7 +106,7 @@ def targs(
     command: str,
     read_from: TextIO,
     shell: bool,
-    template_engine: template.TemplateEngines,
+    template_engine: TemplateExtensions,
     max_procs: int,
     redirect_stdout: bool,
     redirect_stderr: bool,
diff --git a/patternutils/commands/tln.py b/patternutils/commands/tln.py
index f76898d..d125f96 100644
--- a/patternutils/commands/tln.py
+++ b/patternutils/commands/tln.py
@@ -6,6 +6,7 @@ from typing import Any, TextIO
 import click
 
 from .. import config, template, utils
+from ..types import TemplateExtensions
 
 
 def link_file(old_filename: str, new_filename: str, symbolic: bool) -> None:
@@ -106,7 +107,7 @@ def tln(
     editor: bool,
     force: bool,
     symbolic: bool,
-    template_engine: template.TemplateEngines,
+    template_engine: TemplateExtensions,
     use_predefined_template: bool,
     read_from: TextIO,
     key: str,
diff --git a/patternutils/commands/tmv.py b/patternutils/commands/tmv.py
index 150f67d..9198171 100644
--- a/patternutils/commands/tmv.py
+++ b/patternutils/commands/tmv.py
@@ -6,6 +6,7 @@ from typing import Any, TextIO
 import click
 
 from .. import config, template, utils
+from ..types import TemplateExtensions
 
 
 def move_file(old_filename: str, new_filename: str) -> None:
@@ -105,7 +106,7 @@ def tmv(
     interactive: bool,
     editor: bool,
     force: bool,
-    template_engine: template.TemplateEngines,
+    template_engine: TemplateExtensions,
     abort_on_path_exist: bool,
     abort_on_key_not_found: bool,
     use_predefined_template: bool,
diff --git a/patternutils/config.py b/patternutils/config.py
index 7f67f03..b236ec2 100644
--- a/patternutils/config.py
+++ b/patternutils/config.py
@@ -1,34 +1,37 @@
 import glob
 import os.path
 import re
-from typing import Literal, Optional, cast, overload
+from typing import Optional, cast, get_args, overload
 
 import click
 
-TemplateExtensions = Literal["python", "jinja2", "regex"]
-CommandNames = Literal["pmatch", "tmv", "targs", "tln"]
+from .types import (
+    PredefinedCommands,
+    PredefinedExtensions,
+    RegexCommands,
+    RegexExtensions,
+    TemplateCommands,
+    TemplateExtensions,
+)
 
 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"]]: ...
+    template_name: str, command_name: RegexCommands
+) -> tuple[str, RegexExtensions]: ...
 
 
 @overload
 def load_template(
-    template_name: str, command_name: Literal["tmv", "targs", "tln"]
-) -> tuple[str, Literal["python", "jinja2"]]: ...
+    template_name: str, command_name: TemplateCommands
+) -> tuple[str, TemplateExtensions]: ...
 
 
 def load_template(
-    template_name: str, command_name: Optional[CommandNames]
-) -> tuple[str, TemplateExtensions]:
+    template_name: str, command_name: Optional[PredefinedCommands]
+) -> tuple[str, PredefinedExtensions]:
     """Load a template from the config directory.
 
     Args:
@@ -69,7 +72,7 @@ def load_template(
 
 
 def find_templates(
-    template_name: str, command_name: Optional[CommandNames] = None
+    template_name: str, command_name: Optional[PredefinedCommands] = None
 ) -> list[str]:
     """Finds matching template files based on the template name and command.
 
@@ -90,17 +93,28 @@ def find_templates(
         extension = ".*"
     else:
         path_sub_regex = f"(globals|{command_name})"
-        if command_name in REGEX_PROGRAMS:
-            extension = "regex"
-        elif command_name in TEMPLATE_PROGRAMS:
+
+        if command_name in get_args(RegexCommands):
+            extensions = get_args(RegexExtensions)
+        elif command_name in get_args(TemplateCommands):
             # 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)"
+            extensions = get_args(TemplateExtensions)
         else:
             raise ValueError(f"Unknown command '{command_name}'.")
 
+        if len(extensions) > 0:
+            if len(extensions) > 1:
+                extension = f"({'|'.join(extensions)})"
+            else:
+                extension = extensions[0]
+        else:
+            raise ValueError(
+                f"Could not determine file extension for command '{command_name}'."
+            )
+
     pattern = re.compile(
         os.path.join(CONFIG_PATH, path_sub_regex, f"{template_name}.{extension}")
     )
diff --git a/patternutils/template.py b/patternutils/template.py
index c02b687..86e5a97 100644
--- a/patternutils/template.py
+++ b/patternutils/template.py
@@ -1,6 +1,7 @@
-from typing import Callable, Literal, cast
+from typing import Callable, Optional, cast
 
-TemplateEngines = Literal["python", "jinja2"]
+from .config import CONFIG_PATH
+from .types import TemplateCommands, TemplateExtensions
 
 try:
     import jinja2
@@ -52,7 +53,10 @@ else:
 
 
 def get_render_function(
-    template: str, *, engine: TemplateEngines = "python"
+    template: str,
+    *,
+    engine: TemplateExtensions = "python",
+    command_name: Optional[TemplateCommands] = None,
 ) -> Callable[..., str]:
     """Returns a function that renders a template using the specified engine.
 
@@ -68,7 +72,18 @@ def get_render_function(
     """
     if engine == "jinja2":
         if jinja_feature:
-            env_template = jinja_env.from_string(template)
+            if command_name is not None:
+                jinja_env_overlay = jinja_env.overlay(
+                    loader=jinja2.ChoiceLoader(
+                        [
+                            jinja2.FileSystemLoader(CONFIG_PATH / command_name),
+                            jinja2.FileSystemLoader(CONFIG_PATH / "globals"),
+                        ]
+                    )
+                )
+            else:
+                jinja_env_overlay = jinja_env
+            env_template = jinja_env_overlay.from_string(template)
             return cast(Callable[..., str], env_template.render)
         else:
             raise RuntimeError("Jinja2 is required but not installed!")
diff --git a/patternutils/types.py b/patternutils/types.py
new file mode 100644
index 0000000..cf5c8f0
--- /dev/null
+++ b/patternutils/types.py
@@ -0,0 +1,14 @@
+from typing import Literal
+
+PlainCommands = Literal["pmediainfo"]
+
+RegexCommands = Literal["pmatch", "pjson"]
+RegexExtensions = Literal["regex"]
+
+TemplateCommands = Literal["tmv", "targs", "tln"]
+TemplateExtensions = Literal["python", "jinja2"]
+
+PredefinedCommands = RegexCommands | TemplateCommands
+PredefinedExtensions = TemplateExtensions | RegexExtensions
+
+CommandNames = PlainCommands | TemplateCommands | RegexCommands