diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index ad25f3d0c1684dfa55aca802085cb22c44bfa2e3..040e3c2851a3acc3b83c9cefaef8c49a5f9a63d0 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -43,6 +43,7 @@ pip_install:
stage: Install
script:
- pip install .
+ - python -c "import theia_picker"
# ------------------------------- Doc ------------------------------------------
diff --git a/setup.py b/setup.py
index 8f92efa5ed81ecafec05398e941aedbbcdaa3829..9a3c597a8a4f036e51608f741c4850b8deec12e5 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
from setuptools import setup, find_packages
-install_requires = ["requests", "pydantic==1.*", "urllib3", "tqdm"]
+install_requires = ["requests", "pydantic==2.*", "urllib3", "tqdm"]
setup(
name="theia-picker",
diff --git a/test/download_test.py b/test/download_test.py
index 3065a4fcc14b60afecdb6ec19131cb9755bca461..ed67c42640c150d79ca624781e2f8d81b6386553 100644
--- a/test/download_test.py
+++ b/test/download_test.py
@@ -7,6 +7,8 @@ import theia_picker
username = os.environ["THEIA_IDENT"]
password = os.environ["THEIA_PASS"]
+test_zip = False
+
with tempfile.TemporaryDirectory() as output_dir:
# Login
cat = theia_picker.TheiaCatalog(
@@ -23,29 +25,36 @@ with tempfile.TemporaryDirectory() as output_dir:
assert len(features) == 1
print("Search OK")
-# # Download some files
-# patterns = [".jpg", ".xml", "EDG_R2.tif", "FRE_B7.tif"]
-# for feat in features:
-# assert feat.properties.product_identifier == \
-# "SENTINEL2A_20220114-103855-001_L2A_T31TFJ_D"
-# files = feat.list_files_in_archive()
-# assert len(files) == 527
-# for file in files:
-# if any(pattern in file for pattern in patterns):
-# feat.download_single_file(
-# filename=file,
-# download_dir=output_dir
-# )
-# out_file = os.path.join(output_dir, file)
-# assert os.path.isfile(out_file)
-# print("Download single file OK")
-#
-# # Download files batch
-# feats = cat.search(
-# tile_name="T31TEJ", start_date="14/01/2021", level="LEVEL2A"
-# )
-# for f in feats:
-# f.download_files(
-# matching=["FRE_B4.tif", "FRE_B8.tif"], download_dir="/tmp"
-# )
-# print("Download multiple files OK")
+ # Download some files
+ patterns = [".jpg", ".xml", "EDG_R2.tif", "FRE_B7.tif"]
+ for feat in features:
+ assert feat.properties.product_identifier == \
+ "SENTINEL2A_20220114-103855-001_L2A_T31TFJ_D"
+ if test_zip:
+ files = feat.list_files_in_archive()
+ assert len(files) == 527
+ for file in files:
+ if any(pattern in file for pattern in patterns):
+ feat.download_single_file(
+ filename=file,
+ download_dir=output_dir
+ )
+ out_file = os.path.join(output_dir, file)
+ assert os.path.isfile(out_file)
+ print("Download single file OK")
+
+ # Download archive
+ feat.download_archive("/tmp/")
+ print("Download archive OK")
+
+ # Download files batch
+ feats = cat.search(
+ tile_name="T31TEJ", start_date="14/01/2021", level="LEVEL2A"
+ )
+ if test_zip:
+ for f in feats:
+ f.download_files(
+ matching=["FRE_B4.tif", "FRE_B8.tif"], download_dir="/tmp"
+ )
+ print("Download multiple files OK")
+
diff --git a/theia_picker/download.py b/theia_picker/download.py
index 7e12294349afb8d0ef4c99c860ab4f224ca34740..60ad7557d9ad4f8015f1e11010d48752a79448c9 100644
--- a/theia_picker/download.py
+++ b/theia_picker/download.py
@@ -15,7 +15,7 @@ import zlib
from contextlib import nullcontext
from typing import Any, Dict, List, Union, Callable
from urllib.parse import urlencode
-from pydantic import BaseModel, Field, validator, Extra # pylint: disable = no-name-in-module, line-too-long # noqa: E501
+from pydantic import BaseModel, Field, field_validator, FieldValidationInfo, ConfigDict # pylint: disable = no-name-in-module, line-too-long # noqa: E501
from requests.adapters import HTTPAdapter, Retry
import requests
from tqdm.autonotebook import tqdm
@@ -30,7 +30,7 @@ SECONDS_BTW_RETRIES = 2
class ProgressStub:
"""
TQDM stub
-
+
"""
def __init__(self, *args, **kwargs):
"""
@@ -518,8 +518,12 @@ class Download(BaseModel): # pylint: disable = too-few-public-methods
url: str = Field(alias="url")
checksum: str = Field(alias="checksum")
- @validator('url', each_item=True)
- def make_url(cls, url: str) -> str: # pylint: disable=no-self-argument
+ @field_validator("url")
+ def make_url( # pylint: disable=no-self-argument
+ cls,
+ url: str,
+ info: FieldValidationInfo # pylint: disable=unused-argument
+ ) -> str:
"""
Model validator
@@ -557,14 +561,14 @@ class Properties(BaseModel): # pylint: disable = too-few-public-methods
services: Services = Field(alias="services")
-class Feature(BaseModel, extra=Extra.allow):
+class Feature(BaseModel):
"""
Feature model
Extended with custom functions to be helpful
"""
-
- _requests_mgr: RequestsManager
- _remote_zip: RemoteZip = None
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+ requests_mgr: RequestsManager
+ remote_zip: RemoteZip = None
id: str = Field(alias="id")
properties: Properties = Field(alias="properties")
@@ -588,8 +592,8 @@ class Feature(BaseModel, extra=Extra.allow):
"""
if renew_token:
- self._requests_mgr.renew_authorization_headers()
- resp = self._requests_mgr.get(self.properties.services.download.url)
+ self.requests_mgr.renew_authorization_headers()
+ resp = self.requests_mgr.get(self.properties.services.download.url)
try:
tot_size_in_bytes = int(resp.headers.get('content-length', 0))
block_size = 32 * 1024 # 32 Kb
@@ -664,13 +668,13 @@ class Feature(BaseModel, extra=Extra.allow):
"""
if renew_token:
- self._requests_mgr.renew_authorization_headers()
- if not self._remote_zip:
- self._remote_zip = RemoteZip(
+ self.requests_mgr.renew_authorization_headers()
+ if not self.remote_zip:
+ self.remote_zip = RemoteZip(
url=self.properties.services.download.url,
- requests_mgr=self._requests_mgr
+ requests_mgr=self.requests_mgr
)
- return self._remote_zip
+ return self.remote_zip
def list_files_in_archive(self) -> List[str]:
"""
@@ -680,6 +684,7 @@ class Feature(BaseModel, extra=Extra.allow):
List of files in the remote archive
"""
+ print(self.requests_mgr)
return self._get_remote_zip().files_list
@retry(
@@ -781,7 +786,7 @@ class TheiaCatalog: # pylint: disable = too-few-public-methods
# Read THEIA credentials
with open(config_file_json, encoding='UTF-8') as json_file:
credentials = json.load(json_file)
- self._requests_mgr = RequestsManager(credentials=credentials)
+ self.requests_mgr = RequestsManager(credentials=credentials)
self.max_records = max_records
def _query(self, dict_query: dict) -> List[Feature]:
@@ -813,7 +818,7 @@ class TheiaCatalog: # pylint: disable = too-few-public-methods
features = search.json().get("features")
log.debug("Got %s results", len(features))
return [
- Feature(_requests_mgr=self._requests_mgr, **record)
+ Feature(requests_mgr=self.requests_mgr, **record)
for record in features
]