Merge branch 'develop' into main

Author: Calon, Rob

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pyetm"
-version = "1.1.1"
+version = "1.2.0"
 description = "Python-ETM Connector"
 authors = [{name = "Rob Calon", email = "[email protected]"}]
 readme = "README.md"

src/pyetm/client/client.py

@@ -191,8 +191,8 @@ def from_saved_scenario_id(
     def __init__(
         self,
         scenario_id: str | None = None,
-        beta_engine: bool = False,
-        reset: bool = False,
+        engine_url: str | None = None,
+        etm_url: str | None = None,
         token: str | None = None,
         session: RequestsSession | AIOHTTPSession | None = None,
         **kwargs
@@ -204,16 +204,15 @@ def __init__(
         scenario_id : str, default None
             The api_session_id to which the client connects. Can only access
             a limited number of methods when scenario_id is set to None.
-        beta_engine : bool, default False
-            Connect to the beta-engine instead of the production-engine.
-        reset : bool, default False
-            Reset scenario on initalization.
         token : str, default None
             Personal access token to authenticate requests to your
             personal account and scenarios. Detects token automatically
-            from environment when assigned to ETM_ACCESS_TOKEN when
-            connected to production or ETM_BETA_ACCESS_TOKEN when
-            connected to beta.
+            from environment when assigned to ETM_ACCESS_TOKEN.
+        engine_url : str, default None
+            Specify URL that points to ETM engine, default to public engine.
+        etm_url : str, default None
+            Specify URL that points to ETM model (pro), default to public
+            energy transition model.
         session: object instance, default None
             session instance that handles requests to ETM's public API.
             Default to use a RequestsSession.
@@ -238,7 +237,8 @@ def __init__(
         self._session = session
 
         # set engine and token
-        self.beta_engine = beta_engine
+        self.engine_url = engine_url
+        self.etm_url = etm_url
         self.token = token
 
         # set scenario id
@@ -247,10 +247,6 @@ def __init__(
         # set default gqueries
         self.gqueries = []
 
-        # reset scenario on intialization
-        if reset and (scenario_id is not None):
-            self.reset_scenario()
-
         # make message
         msg = (
             "Initialised new Client: "
@@ -282,7 +278,7 @@ def __repr__(self):
         params = {
             **{
                 "scenario_id": self.scenario_id,
-                "beta_engine": self.beta_engine,
+                "engine_url": self.engine_url,
                 "session": self.session
                 },
             **self.__kwargs

src/pyetm/client/scenario.py

@@ -1,9 +1,9 @@
 """Authentication methods"""
 
 from __future__ import annotations
-# from pathlib import Path
 
 import copy
+from urllib.parse import urljoin
 
 import pandas as pd
 from .session import SessionMethods
@@ -87,15 +87,7 @@ def private(self, boolean: bool):
     @property
     def pro_url(self) -> str:
         """get pro url for session id"""
-
-        # specify base url
-        base = 'https://energytransitionmodel.com'
-
-        # update to beta server
-        if self.beta_engine:
-            base = base.replace('https://', 'https://beta.')
-
-        return f'{base}/scenarios/{self.scenario_id}/load'
+        return urljoin(self.etm_url, f'scenarios/{self.scenario_id}/load/')
 
     @property
     def scaling(self):

src/pyetm/client/session.py

@@ -7,6 +7,8 @@
 import copy
 import functools
 
+from urllib.parse import urljoin
+
 import pandas as pd
 
 from pyetm.logger import get_modulelogger
@@ -21,39 +23,62 @@
 class SessionMethods:
     """Core methods for API interaction"""
 
+    @property
+    def _default_engine_url(self) -> str:
+        """default engine url"""
+        return "https://engine.energytransitionmodel.com/api/v3/"
+
+    @property
+    def connected_to_default_engine(self) -> bool:
+        """connected to default engine url?"""
+        return self.engine_url == self._default_engine_url
+
     @property
     def _scenario_header(self) -> dict:
         """get full scenario header"""
         return self._get_scenario_header()
 
     @property
-    def base_url(self) -> str:
-        """"base url for carbon transition model"""
+    def engine_url(self) -> str:
+        """engine URL"""
+        return self._engine_url
 
-        # return beta engine url
-        if self.beta_engine:
-            return "https://beta-engine.energytransitionmodel.com/api/v3/"
+    @engine_url.setter
+    def engine_url(self, url: str | None):
 
-        return "https://engine.energytransitionmodel.com/api/v3/"
+        # default url
+        if url is None:
+            url = self._default_engine_url
+
+        # set engine
+        self._engine_url = str(url)
+
+        # reset token and change base url
+        self.token = None
+        self.session.base_url = self._engine_url
+
+        # reset cache
+        self._reset_cache()
 
     @property
-    def beta_engine(self) -> bool:
-        """connects to beta-engine when True and to production-engine
-        when False."""
-        return self._beta_engine
+    def etm_url(self) -> str:
+        """model URL"""
 
-    @beta_engine.setter
-    def beta_engine(self, boolean: bool) -> None:
-        """set beta engine attribute"""
+        # raise error
+        if self.etm_url is None:
+            raise ValueError("ETModel URL not set on initialisation.")
 
-        # set boolean and reset session
-        self._beta_engine = bool(boolean)
+        return self._etm_url
 
-        # set related settings
-        self.token = None
-        self.session.base_url = self.base_url
+    @etm_url.setter
+    def etm_url(self, url: str | None):
 
-        self._reset_cache()
+        # use default pro location
+        if (url is None) & (self.connected_to_default_engine):
+            url = "https://energytransitionmodel.com/"
+
+        # set etmodel
+        self._etm_url = str(url)
 
     @property
     def scenario_id(self) -> int | None:
@@ -113,14 +138,10 @@ def token(self) -> pd.Series | None:
     @token.setter
     def token(self, token: str | None = None):
 
-        # check environment variables for production token
-        if (token is None) & (not self.beta_engine):
+        # check environment variables for token
+        if token is None:
             token = os.getenv('ETM_ACCESS_TOKEN')
 
-        # check environment variables for beta token
-        if (token is None) & self.beta_engine:
-            token = os.getenv('ETM_BETA_ACCESS_TOKEN')
-
         # store token
         self._token = token
 
@@ -179,11 +200,8 @@ def _get_scenario_header(self):
     def _get_session_id(self, scenario_id: int) -> int:
         """get a session_id for a pro-environment scenario"""
 
-        # make pro url
-        host = "https://energytransitionmodel.com"
-        url = f"{host}/saved_scenarios/{scenario_id}/load"
-
         # extract content from url
+        url = urljoin(self.etm_url, f'saved_scenarios/{scenario_id}/load/')
         content = self.session.request("get", url, decoder='text')
 
         # get session id from content