Staging argument to functions defaulting to config value (#57)

* Add staging arg to all functions, default to config

* Fix requirements

* Update doc

---------

Co-authored-by: Tanguy Fardet <tanguy.fardet@protonmail.com>

bw_hestia_bridge/hestia_api/base_api.py

@@ -34,6 +34,7 @@
 
 def hestia_request(
     endpoint: str,
+    staging: bool,
     query: Optional[dict] = None,
     req_type: Literal["get", "post"] = "get",
 ) -> dict:
@@ -44,14 +45,16 @@ def hestia_request(
     ----------
     endpoint : str
         The API endpoint (e.g. "search").
+    staging : bool
+        Whether to use the staging API.
     query : dict, optional (default: None)
         Additional queries (passed via something like "?q1=v1&q2=v2").
     req_type : str, "get" or "post"
         The type of request that will be performed.
     """
     config = bhb.get_config()
 
-    url = staging_url if config["use_staging"] else stable_url
+    url = staging_url if staging else stable_url
 
     proxies = {
         "http": config["http_proxy"],

bw_hestia_bridge/hestia_api/cycle_graph.py

@@ -1,18 +1,23 @@
-from typing import Union
+from typing import Optional, Union
 
 from ..utils import get_config
 from .base_api import hestia_request
 from .querying import get_hestia_node
 
 
-def get_cycle_graph(cycle_id: str) -> list[dict]:
+def get_cycle_graph(
+    cycle_id: str,
+    staging: Optional[bool] = None,
+) -> list[dict]:
     """
     Get all the cycles that are in the connected graph of `cycle_id`.
 
     Parameters
     ----------
     cycle_id : str
-        Hestia ID of the initial cycle
+        Hestia ID of the initial cycle.
+    staging : bool, optional (default: from configuration)
+        Whether to use the staging API.
 
     Returns
     -------
@@ -38,7 +43,9 @@ def get_cycle_graph(cycle_id: str) -> list[dict]:
         "includeAggregated": True,
     }
 
-    res = hestia_request(f"cycles/{cycle_id}/deep-relations", query=q)
+    staging = staging or get_config("use_staging")
+
+    res = hestia_request(f"cycles/{cycle_id}/deep-relations", staging, query=q)
 
     if not isinstance(res, list):
         api_type = "staging" if get_config("use_staging") else "stable"

bw_hestia_bridge/hestia_api/querying.py

@@ -11,6 +11,7 @@ def search_hestia(
     fields: Optional[list[str]] = None,
     limit: Optional[int] = 10,
     how: Literal["or", "and", "exact"] = "or",
+    staging: Optional[bool] = None,
 ) -> list[dict[str, str]]:
     """
     Search the Hestia database.
@@ -37,6 +38,8 @@ def search_hestia(
         Whether the search tries to match any word in `query` ("or"), all
         words in `query` ("and") or to match the whole query exactly
         ("exact").
+    staging : bool, optional (default: from configuration)
+        Whether to use the staging API.
 
     Returns
     -------
@@ -60,6 +63,8 @@ def search_hestia(
 
     how = how or "or"
 
+    staging = staging or get_config("use_staging")
+
     matches: list[dict] = []
 
     if not isinstance(query, dict):
@@ -103,7 +108,7 @@ def search_hestia(
         "query": {"bool": {"must": matches}},
     }
 
-    res = hestia_request("search", query=q, req_type="post")
+    res = hestia_request("search", staging, query=q, req_type="post")
 
     return res.get("results", [])
 
@@ -112,6 +117,7 @@ def get_hestia_node(
     node_id: Union[str, dict[str, str]],
     node_type: Optional[str] = None,
     data_state: Optional[str] = None,
+    staging: Optional[bool] = None,
 ) -> dict:
     """
     Download the Hestia node associated to `node`.
@@ -132,13 +138,17 @@ def get_hestia_node(
     data_state : str, optional (default: "recalculated")
         Version of the data, by default, use "recalculated" to download the
         more detailed version of the data. Use "original" to get the raw data.
+    staging : bool, optional (default: from configuration)
+        Whether to use the staging API.
 
     Returns
     -------
     node : dict
         The dict associated to the JSON-LD entry describing `node` in the
         Hestia database.
     """
+    staging = staging or get_config("use_staging")
+
     if isinstance(node_id, dict):
         assert "@type" in node_id, "`node` must contain an '@type' entry."
         assert "@id" in node_id, "`node` must contain an '@id' entry."
@@ -153,12 +163,13 @@ def get_hestia_node(
     data_state = data_state or "recalculated"
 
     if node_type == "cycle":
-        return hestia_request(f"{node_type}s/{node_id}?dataState={data_state}")
+        return hestia_request(
+            f"{node_type}s/{node_id}?dataState={data_state}", staging)
     else:
-        return hestia_request(f"{node_type}s/{node_id}")
+        return hestia_request(f"{node_type}s/{node_id}", staging)
 
 
-def get_node_type(node_id: str) -> str:
+def get_node_type(node_id: str, staging: Optional[bool] = None) -> str:
     """
     Get the node type from its Hestia ID
 
@@ -171,16 +182,20 @@ def get_node_type(node_id: str) -> str:
     -------
     node_type : str
         The type of the node.
+    staging : bool, optional (default: from configuration)
+        Whether to use the staging API.
 
     Raises
     ------
     ValueError : if `node_id` is not found.
     """
-    res = search_hestia({"@id": node_id}, how="exact")
+    staging = staging or get_config("use_staging")
+
+    res = search_hestia({"@id": node_id}, how="exact", staging=staging)
 
     if res:
         return res[0]["@type"].lower()
 
-    api_type = "staging" if get_config("use_staging") else "stable"
+    api_type = "staging" if staging else "stable"
 
     raise ValueError(f"The {api_type} API found no node with ID {node_id}.")

bw_hestia_bridge/importer.py

@@ -7,8 +7,8 @@
 from . import set_config
 from .hestia_api import get_cycle_graph, get_hestia_node
 from .strategies import (
+    Converter,
     add_code_from_hestia_attributes,
-    convert,
     drop_zeros,
     link_ecoinvent_biosphere,
     link_ecoinvent_technosphere,
@@ -22,8 +22,8 @@ def __init__(
         ecoinvent_label: str,
         expand_graph: Optional[bool] = True,
         data_state: Literal["original", "recalculated"] = "recalculated",
-        staging: Optional[bool] = False,
         biosphere_label: Optional[str] = "biosphere3",
+        staging: Optional[bool] = None,
     ) -> None:
         """
         Import a Hestia cycle as a Brightway database.
@@ -35,21 +35,24 @@ def __init__(
         data_state : str, optional (default: recalculated)
             Whether to use recalculated data information from Hestia
             or the raw "original" data.
-        staging : bool, optional (default: False)
+        staging : bool, optional (default: from configuration)
             Whether to fetch the cycle from the staging Hestia API.
         """
         # move to staging if necessary
-        set_config("use_staging", staging)
+        self._staging = staging
+        self._convertor = Converter(staging)
 
         # initialize variables
         extended = " expanded" if expand_graph else ""
         self.db_name = f"Hestia cycle {cycle_id}{extended}"
 
         self.cycle_id = cycle_id
 
-        self.data = convert(
-            get_hestia_node(node_id=cycle_id, node_type="cycle", data_state=data_state)
+        self.data = self._convertor.convert(
+            get_hestia_node(node_id=cycle_id, node_type="cycle", data_state=data_state,
+                            staging=staging)
         )
+
         if expand_graph:
             self.get_suppliers(cycle_id=cycle_id, data_state=data_state)
 
@@ -65,15 +68,16 @@ def __init__(
         ]
 
     def get_suppliers(self, cycle_id: str, data_state: str) -> None:
-        graph = get_cycle_graph(cycle_id)
+        graph = get_cycle_graph(cycle_id, staging=self._staging)
         for element in graph:
             if "from" not in element or element["from"].get("@type") != "Cycle":
                 continue
-            nodes = convert(
+            nodes = self._convertor.convert(
                 get_hestia_node(
                     node_id=element["from"]["@id"],
                     node_type="cycle",
                     data_state=data_state,
+                    staging=self._staging,
                 )
             )
             for new_ds in nodes:

bw_hestia_bridge/strategies/__init__.py

@@ -1,4 +1,4 @@
-from .conversion import Converter, convert  # NOQA: F401
+from .conversion import Converter  # NOQA: F401
 from .generic import drop_zeros  # NOQA: F401
 from .linking import (  # NOQA: F401
     add_code_from_hestia_attributes,

bw_hestia_bridge/strategies/conversion.py

@@ -6,15 +6,13 @@
 from ..hestia_api import get_hestia_node
 
 
-def convert(data: dict) -> list:
-    """Convert Hestia data into a BW-compatible form"""
-    return Converter().convert(data)
-
-
 class Converter:
 
     """Convert Hestia data into a BW-compatible form"""
 
+    def __init__(self, staging: Optional[bool] = None):
+        self._staging = staging or get_config("use_staging")
+
     def convert(self, source: dict) -> list:
         return_data = []
 
@@ -69,10 +67,14 @@ def get_basic_metadata(self, obj: dict) -> dict:
 
     @lru_cache
     def get_site(self, node_id: str) -> str:
-        location = get_hestia_node(node_type="site", node_id=node_id)
+        location = get_hestia_node(
+            node_type="site", node_id=node_id, staging=self._staging
+        )
+
         if "name" not in location:
             warnings.warn(f"Can't find location {node_id}; using `GLO` instead")
             return "GLO"
+
         return location
 
     def _add_suffixed(

bw_hestia_bridge/strategies/linking.py

@@ -7,6 +7,7 @@
 from bw2io import activity_hash
 from constructive_geometries import Geomatcher
 
+
 DATA_DIR = Path(__file__).parent.parent.resolve() / "data"
 
 

docs/format.md

@@ -39,7 +39,7 @@ print(transform["products"])
 The supply chain of a cycle is therefore composed of other cycles and
 transformations that are connected via their inputs and outputs, which are
 [``Term``](https://www.hestia.earth/schema/Term) objects in Hestia, that
-will become XXX in Brightway.
+will become reference products in Brightway.
 
 ### Pitfalls and limitations
 

requirements.txt

@@ -2,3 +2,4 @@ bw2data >= 4.0.dev29
 bw2io >= 0.9.DEV23
 platformdirs
 requests
+constructive_geometries