add type hints

1 files changed, 26 insertions, 19 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)