Make load_template more idiomatic

1 files changed, 9 insertions, 7 deletions

diff --git a/patternutils/config.py b/patternutils/config.py
index 4b95cf4..c83d68c 100644
--- a/patternutils/config.py
+++ b/patternutils/config.py
@@ -8,12 +8,13 @@ from typing import Tuple, List, Optional
 CONFIG_PATH = click.get_app_dir('patternutils')
 
 
-def load_template(template_name: str, command_name: str) -> Tuple[str, str]:
+def load_template(template_name: str, command_name: Optional[str]) -> Tuple[str, str]:
     """Load a template from the config directory."""
 
     files = find_templates(template_name, command_name)
     if len(files) == 1:
-        with open(files[0], encoding='utf-8') as f:
+        template_file = files[0]
+        with open(template_file, encoding='utf-8') as f:
             template_lines = f.read()
     elif not files:
         raise FileNotFoundError
@@ -21,13 +22,14 @@ def load_template(template_name: str, command_name: str) -> Tuple[str, str]:
         raise RuntimeError('Multiple files found!')
 
     output_pattern = template_lines.replace('\n', '')
-    template_file = files[0]
-    if template_file.endswith('regex'):
-        template_engine = 'regex'
-    elif template_file.endswith('jinja2'):
-        template_engine = 'jinja2'
+
+    for fileext in ('regex', 'jinja2'):
+        if template_file.endswith(fileext):
+            template_engine = fileext
+            break
     else:
         template_engine = 'python'
+
     return output_pattern, template_engine