feat(copebr): add locales; include Argenina to weather

.gitignore

@@ -54,6 +54,9 @@ htmlcov/
 .coverage
 .coverage.*
 .cache
+.mypy_cache
+.pytest_cache
+.ruff_cache
 nosetests.xml
 coverage.xml
 *satellite_weather_downloader/tests
@@ -97,3 +100,4 @@ rabbitmq.conf
 *certificate.pem
 *key.pem
 *key.p12
+.misc

conda/dev.yaml

@@ -8,4 +8,4 @@ dependencies:
   - make
   - poetry
   - setuptools >= 69.1.0
-  - gdal
+  # - gdal

pyproject.toml

@@ -42,13 +42,14 @@ loguru = "^0.6.0"
 amqp = "^5.1.1"
 requests = "^2.28.2"
 prompt-toolkit = "^3.0.36"
-geopandas = "^0.12.2"
+geopandas = ">=1.0.1"
 matplotlib = "^3.6.3"
 shapely = ">=2.0.3"
 dask = {extras = ["dataframe"], version = "^2024.5.2"}
-xarray = ">=2023.7.0"
-gdal = "^3.8.4"
+xarray = ">=2024.9.0"
+# gdal = "^3.9.3"
 epiweeks = "^2.3.0"
+orjson = "^3.10.7"
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=7.4"

satellite/weather/__init__.py

@@ -4,7 +4,7 @@
 import xarray as xr
 
 from .copebr import *  # noqa
-from .dsei import *  # noqa
+# from .dsei import *  # noqa
 
 
 def load_dataset(file_path: str) -> xr.Dataset:

satellite/weather/copebr.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Union, Literal
 
 from epiweeks import Week
 import dask
@@ -9,7 +9,8 @@
 from loguru import logger  # type: ignore
 from sqlalchemy.engine import Connectable  # type: ignore
 
-from . import brazil
+# from .locales import BR
+from .utils import extract_latlons, extract_coordinates
 
 xr.set_options(keep_attrs=True)
 
@@ -58,8 +59,13 @@ class CopeBRDatasetExtension:
     umid_max   : Maximum┘
     """
 
-    def __init__(self, xarray_ds: xr.Dataset) -> None:
+    def __init__(
+        self,
+        xarray_ds: xr.Dataset,
+        locale: Literal['BR', 'AR'] = 'BR'
+    ) -> None:
         self._ds = xarray_ds
+        self.locale = locale
 
     def to_dataframe(self, geocodes: Union[list, int], raw: bool = False):
         df = _final_dataframe(dataset=self._ds, geocodes=geocodes, raw=raw)
@@ -97,11 +103,15 @@ def to_sql(
             )
             logger.debug(f"{geocode} updated on {schema}.{tablename}")
 
-    def geocode_ds(self, geocode: int, raw: bool = False):
-        return _geocode_ds(self._ds, geocode, raw)
+    def geocode_ds(self, geocode: int | str, raw: bool = False):
+        return _geocode_ds(self._ds, geocode, self.locale, raw)
 
 
-def _final_dataframe(dataset: xr.Dataset, geocodes: Union[list, int], raw=False):
+def _final_dataframe(
+    dataset: xr.Dataset,
+    geocodes: Union[list, int],
+    raw=False
+):
     geocodes = [geocodes] if isinstance(geocodes, int) else geocodes
 
     dfs = []
@@ -167,14 +177,16 @@ def _geocode_to_dataframe(dataset: xr.Dataset, geocode: int, raw=False):
     df = ds.to_dataframe()
     del ds
     geocode = [geocode for g in range(len(df))]
-    df = df.assign(geocodigo=da.from_array(geocode))
+    df = df.assign(geocode=da.from_array(geocode))
     df = df.assign(epiweek=str(Week.fromdate(df.index.to_pydatetime()[0])))
-    columns_to_round = list(set(df.columns).difference(set(["geocodigo", "epiweek"])))
+    columns_to_round = list(set(df.columns).difference(
+        set(["geocode", "epiweek"]))
+    )
     df[columns_to_round] = df[columns_to_round].map(lambda x: np.round(x, 4))
     return df
 
 
-def _geocode_ds(ds: xr.Dataset, geocode: int, raw=False):
+def _geocode_ds(ds: xr.Dataset, geocode: int | str, locale: str, raw=False):
     """
     This is the most important method of the extension. It will
     slice the dataset according to the geocode provided, do the
@@ -193,7 +205,7 @@ def _geocode_ds(ds: xr.Dataset, geocode: int, raw=False):
                     the data corresponds to a 3h interval range for
                     each day in the dataset.
     """
-    lats, lons = _get_latlons(geocode)
+    lats, lons = _get_latlons(geocode, locale)
 
     geocode_ds = _convert_to_br_units(
         _slice_dataset_by_coord(dataset=ds, lats=lats, lons=lons)
@@ -294,13 +306,15 @@ def _reduce_by(ds: xr.Dataset, func, prefix: str):
     )
 
 
-def _get_latlons(geocode: int) -> tuple[list[float], list[float]]:
+def _get_latlons(
+        geocode: int | str, locale: str
+) -> tuple[list[float], list[float]]:
     """
     Extract Latitude and Longitude from a Brazilian's city
     according to IBGE's geocode format.
     """
-    lat, lon = brazil.extract_latlons.from_geocode(int(geocode))
-    N, S, E, W = brazil.extract_coordinates.from_latlon(lat, lon)
+    lat, lon = extract_latlons.from_geocode(int(geocode), locale)
+    N, S, E, W = extract_coordinates.from_latlon(lat, lon)
 
     lats = [N, S]
     lons = [E, W]

satellite/weather/dsei.py

@@ -6,8 +6,6 @@
 from matplotlib.path import Path  # type: ignore
 from shapely.geometry.polygon import Polygon  # type: ignore
 
-from . import brazil  # type: ignore
-
 
 @xr.register_dataset_accessor("DSEI")
 class CopeDSEIDatasetExtension:
@@ -38,7 +36,8 @@ def load_polygons(self):
 
     def get_polygon(self, dsei: Union[str, int]) -> Polygon:
         if self._dsei_df is None:
-            logger.error("Polygons are not loaded. Use `.DSEI.load_poligons()`")
+            logger.error(
+                "Polygons are not loaded. Use `.DSEI.load_poligons()`")
             return None
 
         polygon = self.__do_polygon(dsei)
@@ -49,9 +48,11 @@ def __getitem__(self, __dsei: Union[str, int] = None):
             return self.__do_dataset(__dsei)
         except AttributeError:
             if self._dsei_df is None:
-                logger.error("Polygons are not loaded. Use `.DSEI.load_poligons()`")
+                logger.error(
+                    "Polygons are not loaded. Use `.DSEI.load_poligons()`")
                 return None
-            logger.error(f"{__dsei} not found. List all DSEIs with `.DSEI.info()`")
+            logger.error(
+                f"{__dsei} not found. List all DSEIs with `.DSEI.info()`")
             return None
 
     def __do_grid(self):
@@ -63,7 +64,8 @@ def __do_grid(self):
     def __do_polygon(self, __dsei: Union[str, int]) -> Polygon:
         if isinstance(__dsei, str):
             cod = float(self.DSEIs[self.DSEIs.DSEI == __dsei].code)
-            polygon = self._dsei_df[self._dsei_df.cod_dsei == cod].geometry.item()
+            polygon = self._dsei_df[self._dsei_df.cod_dsei ==
+                                    cod].geometry.item()
         elif isinstance(__dsei, int):
             polygon = self._dsei_df[
                 self._dsei_df.cod_dsei == float(__dsei)