Use unique key as a sub diretory if both a file name and unique key are supplied

Signed-off-by: Bram Stoeller <bram.stoeller@alliander.com>

Author: bramstoeller

src/power_grid_model_io/utils/download.py

@@ -214,20 +214,26 @@ def get_download_path(
         unique_key: A unique string that can be used to generate a filename (e.g. a url).
     """
 
-    # If no file_name is given, generate a file name
-    if file_name is None:
+    # If no dir_path is given, use the system's designated folder for temporary files.
+    if dir_path is None:
+        dir_path = Path(tempfile.gettempdir())
+
+    # If no specific download path was given, we need to generate a unique key (based on the given unique key)
+    if file_name is None or unique_key is not None:
         if unique_key is None:
-            raise ValueError("Supply data in order to auto generate a download path.")
+            raise ValueError("Supply a unique key in order to auto generate a download path.")
 
         sha256 = hashlib.sha256()
         sha256.update(unique_key.encode())
-        hash_str = base64.b64encode(sha256.digest()).decode("ascii")
-        hash_str = hash_str.replace("/", "_").replace("+", "-").rstrip("=")
-        file_name = Path(f"{hash_str}.download")
+        unique_key = base64.b64encode(sha256.digest()).decode("ascii")
+        unique_key = unique_key.replace("/", "_").replace("+", "-").rstrip("=")
 
-    # If no dir_path is given, use the system's designated folder for temporary files
-    elif dir_path is None:
-        dir_path = Path(tempfile.gettempdir())
+        # If no file name was given, use the unique key as a file name
+        if file_name is None:
+            file_name = Path(f"{unique_key}.download")
+        # Otherwise, use the unique key as a sub directory
+        else:
+            dir_path /= unique_key
 
     # Combine the two paths
     assert file_name is not None

tests/unit/utils/test_download.py

@@ -22,6 +22,10 @@
 
 Response = namedtuple("Response", ["status", "headers"])
 
+# The base64 representation of the sha256 hash of "foo" is LCa0a2j/xo/5m0U8HTBBNBNCLXBkg7+g+YpeiGJm564=
+# The / and + will be replaced with a _ and - character and the trailing = character(s) will be removed.
+FOO_KEY = "LCa0a2j_xo_5m0U8HTBBNBNCLXBkg7-g-YpeiGJm564"
+
 
 @pytest.fixture()
 def temp_dir():
@@ -356,25 +360,23 @@ def test_get_response_info__no_length(mock_urlopen):
 
 def test_get_download_path(temp_dir: Path):
     # Act
-    path = get_download_path(dir_path=temp_dir, file_name="file_name.zip", unique_key="foo")
+    path = get_download_path(dir_path=temp_dir, file_name="file_name.zip")
 
     # Assert
     assert path == temp_dir / "file_name.zip"
 
 
 def test_get_download_path__auto_dir():
     # Act
-    path = get_download_path(file_name="file_name.zip")
+    path = get_download_path(file_name="file_name.zip", unique_key="foo")
 
     # Assert
-    assert path == Path(tempfile.gettempdir()).resolve() / "file_name.zip"
+    assert path == Path(tempfile.gettempdir()).resolve() / FOO_KEY / "file_name.zip"
 
 
 def test_get_download_path__auto_file_name(temp_dir: Path):
     # Arrange
-    # The base64 representation of the sha256 hash of "foo" is LCa0a2j/xo/5m0U8HTBBNBNCLXBkg7+g+YpeiGJm564=
-    # The / and + will be replaced with a _ and - character and the trailing = character(s) will be removed.
-    expected_file_name = "LCa0a2j_xo_5m0U8HTBBNBNCLXBkg7-g-YpeiGJm564.download"
+    expected_file_name = f"{FOO_KEY}.download"
 
     # Act
     path = get_download_path(dir_path=temp_dir, unique_key="foo")
@@ -385,7 +387,7 @@ def test_get_download_path__auto_file_name(temp_dir: Path):
 
 def test_get_download_path__missing_data(temp_dir: Path):
     # Act / Assert
-    with pytest.raises(ValueError, match=r"Supply data in order to auto generate a download path\."):
+    with pytest.raises(ValueError, match=r"Supply a unique key in order to auto generate a download path\."):
         get_download_path(dir_path=temp_dir)
 
 

tests/utils.py

@@ -178,6 +178,8 @@ def eq(left, right) -> bool:
                     return False
             if isinstance(left, NDFrame):
                 return (left == right).all()
+            if isinstance(right, NDFrame):
+                return False
             if isnan(left) and isnan(right):
                 return True
             return left == right