Refactor

02_bayern.py

@@ -13,13 +13,12 @@ import random
 import xml.etree.ElementTree as ET
 from pathlib import Path
 
-from citygml_download import CityGMLWithHash, download_all_files, download_file
+from citygml_download import CityGMLWithHash, download_all_files, download_file, SCRIPT_DIR, TMP_DIR
 
-SCRIPT_DIR = Path(__file__).resolve().parent
 BAYERN_METALINK = "https://geodaten.bayern.de/odd/a/lod2/citygml/meta/metalink/09.meta4"
 
 
-def download_metalink(metalink_url: str, tmp_path: Path) -> Path:
+def download_metalink(metalink_url: str, tmp_path: Path = TMP_DIR) -> Path:
     """Download metalink file to temporary directory."""
     print(f"Downloading {metalink_url}")
     basename = metalink_url.split("/")[-1] or "metalink.xml"
@@ -65,14 +64,12 @@ def parse_metalink(metalink_path: Path, download_dir: Path) -> list[CityGMLWithH
 
 if __name__ == "__main__":
     # Download and parse metalink
-    tmp_dir = SCRIPT_DIR / "tmp"
     download_dir = SCRIPT_DIR / "citygml" / "bayern"
 
-    tmp_dir.mkdir(parents=True, exist_ok=True)
     download_dir.mkdir(parents=True, exist_ok=True)
 
-    metalink_path = download_metalink(BAYERN_METALINK, tmp_dir)
-    files = parse_metalink(metalink_path, download_dir)
+    metalink_path = download_metalink(BAYERN_METALINK)
+    citygmls = parse_metalink(metalink_path, download_dir)
 
     # Download all files
-    download_all_files(files, tmp_dir, sleep=5, retries=1, jobs=1)
+    download_all_files(citygmls)

citygml_download.py

@@ -17,14 +17,14 @@ from urllib.error import URLError
 from urllib.request import urlopen
 
 SCRIPT_DIR = Path(__file__).resolve().parent
+TMP_DIR = SCRIPT_DIR / "tmp"
+TMP_DIR.mkdir(parents=True, exist_ok=True)
 
 
 def download_all_files(
-    files: Sequence["CityGML"], tmp_dir: Path = SCRIPT_DIR / "tmp", jobs: int = 3, retries: int = 2, sleep: int = 5
+    files: Sequence["CityGML"], tmp_dir: Path = TMP_DIR, jobs: int = 3, retries: int = 2, sleep: int = 5
 ):
     """Download all CityGML files with parallel downloads."""
-    # Create directories
-    tmp_dir.mkdir(parents=True, exist_ok=True)
 
     print(f"Found {len(files)} files to download\n")
 
@@ -141,7 +141,7 @@ class CityGML(ABC):
 
         return self.verify(self.path)
 
-    def download(self, tmp_dir: Path, max_retries: int = 3, sleep: int = 0) -> bool:
+    def download(self, tmp_dir: Path = TMP_DIR, max_retries: int = 3, sleep: int = 0) -> bool:
         """
         Download and verify the file.