Split download en zip_files util

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

Author: bramstoeller

src/power_grid_model_io/utils/zip_files.py

@@ -0,0 +1,79 @@
+# SPDX-FileCopyrightText: 2022 Contributors to the Power Grid Model IO project <dynamic.grid.calculation@alliander.com>
+#
+# SPDX-License-Identifier: MPL-2.0
+"""
+Helper function to extract zip files
+
+csv_dir_path = extract("/tmp/1-complete_data-mixed-all-0-sw.zip")
+
+This extracts the files, in a folder which corresponds to the zip file name ("/tmp/1-complete_data-mixed-all-0-sw/" in
+our example), and it returns the path to that directory. By default, it will not re-download or re-extract the zip
+file as long as the files exist.
+
+"""
+
+import zipfile
+from pathlib import Path
+from typing import Optional
+
+import structlog
+from tqdm import tqdm
+
+_log = structlog.get_logger(__name__)
+
+
+def extract(src_file_path: Path, dst_dir_path: Optional[Path] = None, skip_if_exists=False) -> Path:
+    """
+    Extract a .zip file and return the destination dir
+
+    Args:
+        src_file_path: The .zip file to extract.
+        dst_dir_path: An optional destination path. If none is given, the src_file_path without .zip extension is used.
+        skip_if_exists: Skip existing files, otherwise raise an exception when a file exists.
+
+    Returns: The path where the files are extracted
+
+    """
+    if src_file_path.suffix.lower() != ".zip":
+        raise ValueError(f"Only files with .zip extension are supported, got {src_file_path.name}")
+
+    if dst_dir_path is None:
+        dst_dir_path = src_file_path.with_suffix("")
+
+    log = _log.bind(src_file_path=src_file_path, dst_dir_path=dst_dir_path)
+
+    if dst_dir_path.exists():
+        if not dst_dir_path.is_dir():
+            raise NotADirectoryError(f"Destination dir {dst_dir_path} exists and is not a directory")
+
+    # Create the destination directory
+    dst_dir_path.mkdir(parents=True, exist_ok=True)
+
+    # Extract per file, so we can show a progress bar
+    with zipfile.ZipFile(src_file_path, "r") as zip_file:
+        file_list = zip_file.namelist()
+        for file_path in tqdm(desc="Extracting", iterable=file_list, total=len(file_list), unit="file", leave=True):
+            dst_file_path = dst_dir_path / file_path
+            if dst_file_path.exists() and dst_file_path.stat().st_size > 0:
+                if skip_if_exists:
+                    log.debug("Skip file extraction, destination file exists", dst_file_path=dst_file_path)
+                    continue
+                raise FileExistsError(f"Destination file {dst_dir_path / file_path} exists and is not empty")
+            zip_file.extract(member=file_path, path=dst_dir_path)
+
+    # Zip files often contain a single directory with the same name as the zip file.
+    # In that case, return the dir to that directory instead of the root dir
+    only_item: Optional[Path] = None
+    for item in dst_dir_path.iterdir():
+        # If only_item is None, this is the first iteration, so item may be the only item
+        if only_item is None:
+            only_item = item
+        # Else, if only_item is not None, there are more than one items in the root of the directory.
+        # This means hat there is no 'only_item' and we can stop the loop
+        else:
+            only_item = None
+            break
+    if only_item and only_item.is_dir() and only_item.name == src_file_path.stem:
+        dst_dir_path = only_item
+
+    return dst_dir_path.resolve()

tests/unit/utils/test_zip_files.py

@@ -0,0 +1,115 @@
+# SPDX-FileCopyrightText: 2022 Contributors to the Power Grid Model IO project <dynamic.grid.calculation@alliander.com>
+#
+# SPDX-License-Identifier: MPL-2.0
+import shutil
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import structlog.testing
+
+from power_grid_model_io.utils.zip_files import extract
+
+from ...utils import MockTqdm, assert_log_exists
+
+DATA_DIR = Path(__file__).parents[2] / "data" / "zip_files"
+ZIP1 = DATA_DIR / "foo.zip"
+ZIP2 = DATA_DIR / "foo-bar.zip"
+
+
+@pytest.fixture()
+def temp_dir():
+    with tempfile.TemporaryDirectory() as tmp:
+        yield Path(tmp)
+
+
+@patch("power_grid_model_io.utils.download.tqdm")
+def test_extract(mock_tqdm: MagicMock, temp_dir: Path):
+    # Arrange
+    src_file_path = temp_dir / "compressed.zip"
+    dst_dir_path = temp_dir / "extracted"
+    shutil.copyfile(ZIP2, src_file_path)
+    mock_tqdm.side_effect = MockTqdm
+
+    # Act
+    extract_dir_path = extract(src_file_path=src_file_path, dst_dir_path=dst_dir_path)
+
+    # Assert
+    assert extract_dir_path == dst_dir_path
+    assert (dst_dir_path / "foo.txt").is_file()
+    assert (dst_dir_path / "folder/bar.txt").is_file()
+
+
+@patch("power_grid_model_io.utils.download.tqdm")
+def test_extract__auto_dir(mock_tqdm: MagicMock, temp_dir: Path):
+    # Arrange
+    src_file_path = temp_dir / "compressed.zip"
+    shutil.copyfile(ZIP2, src_file_path)
+    mock_tqdm.side_effect = MockTqdm
+
+    # Act
+    extract_dir_path = extract(src_file_path=src_file_path)
+
+    # Assert
+    assert extract_dir_path == temp_dir / "compressed"
+    assert (temp_dir / "compressed" / "foo.txt").is_file()
+    assert (temp_dir / "compressed" / "folder" / "bar.txt").is_file()
+
+
+def test_extract__invalid_file_extension():
+    # Act / Assert
+    with pytest.raises(ValueError, match=r"Only files with \.zip extension are supported, got tempfile\.download"):
+        extract(src_file_path=Path("/tmp/dir/tempfile.download"))
+
+
+@patch("power_grid_model_io.utils.download.tqdm")
+def test_extract__file_exists(mock_tqdm: MagicMock, temp_dir: Path):
+    # Arrange
+    src_file_path = temp_dir / "compressed.zip"
+    dst_dir_path = temp_dir / "extracted"
+    shutil.copyfile(ZIP2, src_file_path)
+    mock_tqdm.side_effect = MockTqdm
+
+    dst_dir_path.mkdir()
+    with open(dst_dir_path / "foo.txt", "wb") as fp:
+        fp.write(b"\0")
+
+    # Act / Assert
+    with pytest.raises(FileExistsError, match=r"Destination file .*foo\.txt exists and is not empty"):
+        extract(src_file_path=src_file_path, dst_dir_path=dst_dir_path)
+
+
+@patch("power_grid_model_io.utils.download.tqdm")
+def test_extract__skip_if_exists(mock_tqdm: MagicMock, temp_dir: Path):
+    # Arrange
+    src_file_path = temp_dir / "compressed.zip"
+    dst_dir_path = temp_dir / "compressed"
+    shutil.copyfile(ZIP2, src_file_path)
+    mock_tqdm.side_effect = MockTqdm
+
+    dst_dir_path.mkdir()
+    with open(dst_dir_path / "foo.txt", "wb") as fp:
+        fp.write(b"\0")
+
+    # Act / Assert
+    with structlog.testing.capture_logs() as capture:
+        extract(src_file_path=src_file_path, dst_dir_path=dst_dir_path, skip_if_exists=True)
+        assert_log_exists(
+            capture, "debug", "Skip file extraction, destination file exists", dst_file_path=dst_dir_path / "foo.txt"
+        )
+
+
+@patch("power_grid_model_io.utils.download.tqdm")
+def test_extract__return_subdir_path(mock_tqdm: MagicMock, temp_dir: Path):
+    # Arrange
+    src_file_path = temp_dir / "foo.zip"
+    shutil.copyfile(ZIP1, src_file_path)
+    mock_tqdm.side_effect = MockTqdm
+
+    # Act
+    extract_dir_path = extract(src_file_path=src_file_path)
+
+    # Assert
+    assert extract_dir_path == temp_dir / "foo" / "foo"
+    assert (temp_dir / "foo" / "foo" / "foo.txt").is_file()