add type hints

2 files changed, 30 insertions, 25 deletions

diff --git a/src/pymdvar/pymdvar.py b/src/pymdvar/pymdvar.py
index b6ea344..d81f785 100644
--- a/src/pymdvar/pymdvar.py
+++ b/src/pymdvar/pymdvar.py
@@ -1,44 +1,51 @@
 from os import getenv
+from re import Match
+from typing import Any
+from xml.etree.ElementTree import Element
+from markdown import Markdown
 from markdown.extensions import Extension
 from markdown.inlinepatterns import Pattern
 
-VARIABLE_RE = r'(\$\{)([a-zA-Z_]*)(\})'
+VAR_RE: str = r'(\$\{)([a-zA-Z_]*)(\})'
 
 
-class VariablePattern(Pattern):
+class VarPattern(Pattern):
     # need to redefine as an extra attribute needs to be passed
-    def __init__(self, pattern, variables, enable_env, md=None):
-        self.variables = variables
-        self.enable_env = enable_env
+    def __init__(self, pattern: Any,
+                 vars: dict[str, str],
+                 enable_env: bool,
+                 md: Markdown | None = None) -> None:
+        self.vars: dict[str, str] = vars
+        self.enable_env: bool = enable_env
         super().__init__(pattern, md)
 
-    def handleMatch(self, m):
+    def handleMatch(self, m: Match[str]) -> str | Element | None:
         # for some reason the group is offest by 1
-        variable = m.group(3)
-        value = ''
+        var: str | Any = m.group(3)
+        value: str = ''
 
-        if variable in self.variables:
-            value = self.variables[variable]
+        if var in self.vars:
+            value = self.vars[var]
         else:
             if self.enable_env:
-                value = getenv(variable, '')
+                value = getenv(var, '')
         return value
 
 
 class VariableExtension(Extension):
-    def __init__(self, **kwargs):
-        self.config = {
+    def __init__(self, **kwargs: Any) -> None:
+        self.config: dict[str, list[Any | str]] = {
             'enable_env': [False, 'Enable environment variables parsing.'],
             'variables': [dict(), 'Dictionary holding variables to be used.']
         }
         super().__init__(**kwargs)
 
-    def extendMarkdown(self, md):
-        variables = self.getConfig('variables', dict())
-        enable_env = self.getConfig('enable_env', False)
-        variable_pattern = VariablePattern(VARIABLE_RE, variables, enable_env)
-        md.inlinePatterns.register(variable_pattern, 'variable', 175)
+    def extendMarkdown(self, md: Markdown) -> None:
+        vars: dict[str, str] | Any = self.getConfig('variables', dict())
+        enable_env: bool = self.getConfig('enable_env', False)
+        var_pattern: VarPattern = VarPattern(VAR_RE, vars, enable_env)
+        md.inlinePatterns.register(var_pattern, 'variable', 175)
 
 
-def makeExtension(*args, **kwargs):
+def makeExtension(*args: Any, **kwargs: Any):
     return VariableExtension(*args, **kwargs)
diff --git a/tests/test_pymdvar.py b/tests/test_pymdvar.py
index 69ffe56..91a97d0 100644
--- a/tests/test_pymdvar.py
+++ b/tests/test_pymdvar.py
@@ -4,10 +4,8 @@ from pymdvar import VariableExtension
 
 
 def test_empty_input():
-    in_str = ''
-    out_str = markdown(in_str, extensions=[VariableExtension()])
-    print(in_str)
-    print(out_str)
+    in_str: str = ''
+    out_str: str = markdown(in_str, extensions=[VariableExtension()])
     assert in_str == out_str
 
 
@@ -20,7 +18,7 @@ def test_empty_input():
     ('foo **$test}** bar', '<p>foo <strong>$test}</strong> bar</p>'),
 ])
 def test_non_replacements(in_str, exp_str):
-    out_str = markdown(in_str, extensions=[VariableExtension()])
+    out_str: str = markdown(in_str, extensions=[VariableExtension()])
     assert out_str == exp_str
 
 
@@ -32,6 +30,6 @@ def test_non_replacements(in_str, exp_str):
     ('foo ![image](${test}/a.jpg) bar', '<p>foo <img alt="image" src="value/a.jpg" /> bar</p>'),
 ])
 def test_simple_replacements(in_str, exp_str):
-    out_str = markdown(in_str, extensions=[VariableExtension(variables={'test':'value'})])
+    out_str: str = markdown(in_str, extensions=[VariableExtension(variables={'test':'value'})])
     assert out_str == exp_str