Add tests, style checks, update strict_types

9 files changed, 200 insertions, 45 deletions

diff --git a/.gitignore b/.gitignore
index fdf69b4..271895e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,3 +3,4 @@ dist
 *.egg-info
 __pycache__
 *.pyc
+.tox
diff --git a/requirements.txt b/requirements.txt
index e69de29..541e100 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -0,0 +1,3 @@
+pytest
+pylint
+tox
\ No newline at end of file
diff --git a/tests/test_functions.py b/tests/test_functions.py
new file mode 100644
index 0000000..d71f616
--- /dev/null
+++ b/tests/test_functions.py
@@ -0,0 +1,51 @@
+import os
+from tempfile import TemporaryDirectory
+import pytest
+
+from unixreg.constants import *
+from unixreg.functions import *
+
+
+_temp_dir = TemporaryDirectory()
+os.environ["XDG_CONFIG_HOME"] = os.path.join(_temp_dir.name, "unixreg")
+
+def test_ConnectRegistry():
+    key = ConnectRegistry(None, HKEY_CURRENT_USER)
+    key.Close()
+
+    with pytest.raises(OSError):
+        key = ConnectRegistry("FakeComputer", HKEY_CURRENT_USER)
+        key.Close()
+
+def test_CreateKey():
+    def _create(key, subkey):
+        with CreateKey(key, subkey) as key:
+            CloseKey(key)
+
+    _create(HKEY_CURRENT_USER, None)
+    _create(HKEY_CURRENT_USER, "str")
+    _create(HKEY_CURRENT_USER, RegKey("RegKey"))
+
+    with pytest.raises(TypeError):
+        _create(HKEY_CURRENT_USER, 1)
+        _create(HKEY_CURRENT_USER, [])
+        _create(HKEY_CURRENT_USER, ())
+        _create(HKEY_CURRENT_USER, {})
+
+def test_QueryValue():
+    value = "value"
+    value_name = "value_name"
+
+    print(__import__("unixreg").HKEY_CURRENT_USER,)
+    key = CreateKey(HKEY_CURRENT_USER, None)
+    SetValue(key, value_name, REG_SZ, value)
+    assert QueryValue(key, value_name) == value
+
+
+def test_ExpandEnvironmentStrings():
+    _test_str = "testvar"
+    os.environ["TEST"] = _test_str
+    os.environ["HOME"] = _test_str
+
+    assert ExpandEnvironmentStrings(r"%TEST%") == os.environ["TEST"]
+    assert ExpandEnvironmentStrings(r"%USERPROFILE%") == os.environ["HOME"]
\ No newline at end of file
diff --git a/tox.ini b/tox.ini
new file mode 100644
index 0000000..d100e60
--- /dev/null
+++ b/tox.ini
@@ -0,0 +1,8 @@
+[tox]
+envlist = py38
+
+[testenv]
+deps = pytest
+commands = pytest
+setenv =
+    TOX = 1
\ No newline at end of file
diff --git a/unixreg/__init__.py b/unixreg/__init__.py
index 974ecec..bf99ffd 100644
--- a/unixreg/__init__.py
+++ b/unixreg/__init__.py
@@ -1,3 +1,7 @@
+"""
+__init__
+"""
+
 __version__ = "0.1.1"
 
 
diff --git a/unixreg/constants.py b/unixreg/constants.py
index 4a1e4b9..e969c5b 100644
--- a/unixreg/constants.py
+++ b/unixreg/constants.py
@@ -1,13 +1,21 @@
+"""
+Constant Values
+
+https://docs.python.org/3/library/winreg.html#constants
+"""
+
+from .key import RegKey
+
 # HKEY_*
-HKEY_CLASSES_ROOT = "HKEY_CLASSES_ROOT"
-HKEY_CURRENT_USER = "HKEY_CURRENT_USER"
-HKEY_LOCAL_MACHINE = "HKEY_LOCAL_MACHINE"
-HKEY_USERS = "HKEY_USERS"
-HKEY_PERFORMANCE_DATA = "HKEY_PERFORMANCE_DATA"
-HKEY_CURRENT_CONFIG = "HKEY_CURRENT_CONFIG"
-HKEY_DYN_DATA = "HKEY_DYN_DATA"
-
-# https://docs.microsoft.com/en-us/windows/win32/secauthz/standard-access-rights
+HKEY_CLASSES_ROOT = RegKey("HKEY_CLASSES_ROOT")
+HKEY_CURRENT_USER = RegKey("HKEY_CURRENT_USER")
+HKEY_LOCAL_MACHINE = RegKey("HKEY_LOCAL_MACHINE")
+HKEY_USERS = RegKey("HKEY_USERS")
+HKEY_PERFORMANCE_DATA = RegKey("HKEY_PERFORMANCE_DATA")
+HKEY_CURRENT_CONFIG = RegKey("HKEY_CURRENT_CONFIG")
+HKEY_DYN_DATA = RegKey("HKEY_DYN_DATA")
+
+#https://docs.microsoft.com/en-us/windows/win32/secauthz/standard-access-rights
 STANDARD_RIGHTS_REQUIRED = 1
 STANDARD_RIGHTS_WRITE = 1
 STANDARD_RIGHTS_READ = 1
@@ -23,7 +31,9 @@ KEY_CREATE_LINK = 1 << 6
 
 KEY_WRITE = STANDARD_RIGHTS_WRITE ^ KEY_SET_VALUE ^ KEY_CREATE_SUB_KEY
 KEY_READ = STANDARD_RIGHTS_READ ^ KEY_QUERY_VALUE ^ KEY_ENUMERATE_SUB_KEYS
-KEY_ALL_ACCESS = STANDARD_RIGHTS_REQUIRED ^ KEY_QUERY_VALUE ^ KEY_SET_VALUE ^ KEY_CREATE_SUB_KEY ^ KEY_ENUMERATE_SUB_KEYS ^ KEY_NOTIFY ^ KEY_CREATE_LINK
+KEY_ALL_ACCESS = (STANDARD_RIGHTS_REQUIRED ^ KEY_QUERY_VALUE ^ KEY_SET_VALUE
+                 ^ KEY_CREATE_SUB_KEY ^ KEY_ENUMERATE_SUB_KEYS ^ KEY_NOTIFY
+                 ^ KEY_CREATE_LINK)
 
 # 64-bit Specific
 KEY_WOW64_64KEY = 1 << 7
@@ -44,4 +54,4 @@ REG_RESOURCE_LIST = 1 << 10
 REG_FULL_RESOURCE_DESCRIPTOR = 1 << 11
 REG_RESOURCE_REQUIREMENTS_LIST = 1 << 12
 REG_SZ = 1 << 13
-
+ 
\ No newline at end of file
diff --git a/unixreg/functions.py b/unixreg/functions.py
index 200a720..2bd544d 100644
--- a/unixreg/functions.py
+++ b/unixreg/functions.py
@@ -1,39 +1,61 @@
+# pylint: disable=invalid-name, unused-argument, unspecified-encoding, missing-function-docstring
+"""
+Implements all winreg functions
+
+https://docs.python.org/3/library/winreg.html#functions
+"""
 import os
 from typing import Union
+from re import findall
+from tempfile import TemporaryDirectory
+from warnings import warn
 
 from .key import RegKey
 from .constants import STANDARD_RIGHTS_REQUIRED, KEY_WOW64_64KEY, KEY_WRITE, KEY_READ
 from .utils import strict_types
 
 KEY_TYPE = Union[str, RegKey]
-SUBKEY_TYPE = KEY_TYPE | Union[None]
+SUBKEY_TYPE = Union[str, RegKey, None]
 
 _KEY_CACHE = []
 _ENV_REPLACE = {
-    "USERPROFILE": os.getenv("HOME")
+    "USERPROFILE": "HOME"
 }
 
 _CONFIG_DIR = os.getenv("XDG_CONFIG_HOME")
 if not _CONFIG_DIR:
-    _CONFIG_DIR = os.path.join(os.getenv("HOME"), ".config")
+    home = os.getenv("HOME")
+    if home:
+        _CONFIG_DIR = os.path.join(home, ".config")
+    else:
+        _CONFIG_DIR = TemporaryDirectory().name
+        if not os.getenv("TOX"):
+            warn(f"Could not find directory to put registry in. Falling back to {_CONFIG_DIR}")
 _CONFIG_DIR = os.path.join(_CONFIG_DIR, "unixreg")
 
+@strict_types
 def __init_values(key: KEY_TYPE, sub_key: SUBKEY_TYPE = None, access = STANDARD_RIGHTS_REQUIRED):
     if isinstance(key, str):
         key = RegKey(key)
 
-    if sub_key:
+    if sub_key is not None:
+        print(sub_key)
         key = key + sub_key
     key.access = access
 
     return key
 
+@strict_types
 def __create_key(key: RegKey):
     path = os.path.join(_CONFIG_DIR, key.key)
 
     os.makedirs(path, exist_ok=True)
 
 def CloseKey(key: RegKey):
+    """
+    Closes a previously opened registry key.
+    The key argument specifies a previously opened key.
+    """
     key.Close()
 
     try:
@@ -41,12 +63,15 @@ def CloseKey(key: RegKey):
     except ValueError:
         pass
 
-def ConnectRegistry(computer: SUBKEY_TYPE, key: str):
+def ConnectRegistry(computer: Union[str, None], key: RegKey):
+    """
+    Opens a registry handle on another computer and returns the handle
+
+    If computer_name is None, the local computer is used, otherwise
+    OSError is raised to signify the function failing
+    """
     if not computer:
         return OpenKey(key, None)
-
-    # ConnectRegistry is expected to throw an OSError on failure
-    # any program that fails to catch this is to blame
     raise OSError("Not Implemented")
 
 @strict_types
@@ -96,8 +121,16 @@ def EnumValue(key: KEY_TYPE, index: int):
     raise NotImplementedError("Not Implemented")
 
 def ExpandEnvironmentStrings(env: str):
-    for var in _ENV_REPLACE:
-        env = env.replace(f"%{var}%", _ENV_REPLACE[var])
+    for key, val in _ENV_REPLACE.items():
+        env = env.replace(f"%{key}%", f"%{val}%")
+
+    match = findall(r"%(.+?)%", env)
+
+    for val in match:
+        valenv = os.getenv(val)
+        if valenv:
+            env = env.replace(f"%{val}%", valenv)
+
     env.replace("\\", os.path.sep)
     return env
 
@@ -127,10 +160,13 @@ def SaveKey(key: RegKey, file_name: str) -> None:
     # this requires a win32 permission compatibility layer
     raise OSError("Not Implemented")
 
-def SetValue(key: KEY_TYPE, sub_key: str, type: int, value: str):
-    return SetValueEx(key, sub_key, 0, type, value)
+def SetValue(key: KEY_TYPE, sub_key: SUBKEY_TYPE, typearg: int, value: str):
+    if isinstance(sub_key, RegKey):
+        sub_key = sub_key.key
+
+    return SetValueEx(key, sub_key, 0, typearg, value)
 
-def SetValueEx(key: KEY_TYPE, value_name: str, reserved: int, type: int, value: str) -> None:
+def SetValueEx(key: KEY_TYPE, value_name: str, reserved: int, typearg: int, value: str) -> None:
     key = __init_values(key)
 
     filepath = os.path.join(_CONFIG_DIR, key.key, value_name)
@@ -148,7 +184,20 @@ def QueryReflectionKey(key: KEY_TYPE):
 
 
 # Non winreg functions
-def LoadRegFile(file_name: str): str
+def LoadRegFile(file_name: str) -> str:
+
+    def _strip_quotes(val) -> str:
+        _QUOTE_LIST = ("\"", '\'')
+        if val.startswith(_QUOTE_LIST) and val.endswith(_QUOTE_LIST):
+            val = val[1:-1]
+        return val
+
+    def _strip_brackets(val) -> str:
+        _BRACKET_LIST = ("[", "]")
+        if val.startswith(_BRACKET_LIST) and val.endswith(_BRACKET_LIST):
+            val = val[1:-1]
+        return val
+
     with open(file_name, "r") as reg:
         nextline = reg.readline()
 
@@ -158,20 +207,22 @@ def LoadRegFile(file_name: str): str
             line = nextline.strip()
             nextline = reg.readline()
 
-            if len(line) == 1: continue
+            if len(line) == 1:
+                continue
+
             split = line.split("=")
 
-            keyline = strip_brackets(line)
+            keyline = _strip_brackets(line)
             if keyline:
                 key = keyline
             elif key and len(split) == 2:
                 name, value = split
-                name = strip_quotes(name)
-                value = strip_quotes(value)
+                name = _strip_quotes(name)
+                value = _strip_quotes(value)
 
                 os.makedirs(key, exist_ok=True)
 
                 with open(os.path.join(_CONFIG_DIR, key.key, name), "w") as regvalue:
                     regvalue.write(value)
 
-                print(f"[{key}] {name}={value}")
\ No newline at end of file
+                print(f"[{key}] {name}={value}")
diff --git a/unixreg/key.py b/unixreg/key.py
index c4d5134..eccf9b6 100644
--- a/unixreg/key.py
+++ b/unixreg/key.py
@@ -1,16 +1,26 @@
+# pylint: disable=invalid-name, global-statement
+"""
+Implements anything related to the Registry Handle
+"""
 import os
 from copy import deepcopy
 from typing import TypeVar, Union
 
-from .constants import STANDARD_RIGHTS_REQUIRED
-
 RegKeyT = TypeVar("RegKeyT", bound="RegKey")
 
 _HANDLE_COUNTER = 0
 
 class RegKey:
+    """
+    Implementation of the Registry Handle
+    https://docs.python.org/3/library/winreg.html#registry-handle-objects
+    """
+
+    key = ""
+    handle = 0
+    access = 0
 
-    def __init__(self, key: str = "", access: int = STANDARD_RIGHTS_REQUIRED):
+    def __init__(self, key: str = "", access: int = 1):
         global _HANDLE_COUNTER
         _HANDLE_COUNTER += 1
 
@@ -28,25 +38,32 @@ class RegKey:
             retval.key = os.path.join(self.key, other)
             return retval
 
-        return None
+        raise TypeError("Invalid Type")
 
     def __enter__(self) -> RegKeyT:
         return self
 
     def __exit__(self, *args, **kwargs):
-        pass
+        self.Close()
 
     def __repr__(self):
         return __class__.__name__
 
     def __str__(self):
-        return f"{__class__.__name__}({self.key}, {self.handle}, {self.access})"
+        return f"{__class__.__name__}(\"{self.key}\", {self.handle}, {self.access})"
 
     def Close(self):
-        pass
+        """
+        Closes the key by cleaning up its values
+        """
 
-    def Detach(self):
-        pass
+    def Detach(self) -> int:
+        """
+        Is suppose to detach the handle and give it to someone else,
+        but we don't have that luxury
+        """
+        self.Close()
+        return self.handle
 
 
-PyHKEY = RegKey
\ No newline at end of file
+PyHKEY = RegKey
diff --git a/unixreg/utils.py b/unixreg/utils.py
index a5dcf45..9435ceb 100644
--- a/unixreg/utils.py
+++ b/unixreg/utils.py
@@ -1,13 +1,23 @@
+"""
+Utilities
+"""
+from typing import get_args
 
 def strict_types(function):
+    """
+    enforce type hinting on functions
+    """
+
     def _decorator(*args, **kwargs):
         hints = function.__annotations__
         all_args = kwargs.copy()
         all_args.update(dict(zip(function.__code__.co_varnames, args)))
         for argument, argument_type in [(i, type(j)) for i, j in all_args.items()]:
-            if argument in hints:
-                if argument_type != hints[argument]:
-                    raise TypeError('{} is not {}'.format(argument, hints[argument].__name__))
-        result = function(*args, **kwargs)
-        return result
+            if (argument in hints and
+               argument_type != hints[argument]
+               and not isinstance(argument, get_args(hints[argument]))):
+                raise TypeError(f"{argument} is not {hints[argument]}")
+
+        return function(*args, **kwargs)
+
     return _decorator