diff --git a/scaleapi/__init__.py b/scaleapi/__init__.py index 23e4a21..37492b3 100644 --- a/scaleapi/__init__.py +++ b/scaleapi/__init__.py @@ -44,8 +44,10 @@ class Batchlist(Paginator[Batch]): class ScaleClient: """Main class serves as an interface for Scale API""" - def __init__(self, api_key, source=None): - self.api = Api(api_key, source) + def __init__(self, api_key, source=None, api_instance_url=None): + self.api = Api( + api_key, user_agent_extension=source, api_instance_url=api_instance_url + ) def get_task(self, task_id: str) -> Task: """Fetches a task. diff --git a/scaleapi/api.py b/scaleapi/api.py index 1319c8a..0f61199 100644 --- a/scaleapi/api.py +++ b/scaleapi/api.py @@ -7,7 +7,7 @@ from ._version import __package_name__, __version__ from .exceptions import ExceptionMap, ScaleException -SCALE_ENDPOINT = "https://api.scale.com/v1" +SCALE_API_BASE_URL_V1 = "https://api.scale.com/v1" # Parameters for HTTP retry HTTP_TOTAL_RETRIES = 3 # Number of total retries @@ -19,7 +19,7 @@ class Api: """Internal Api reference for handling http operations""" - def __init__(self, api_key, user_agent_extension=None): + def __init__(self, api_key, user_agent_extension=None, api_instance_url=None): if api_key == "" or api_key is None: raise ScaleException("Please provide a valid API Key.") @@ -33,6 +33,7 @@ def __init__(self, api_key, user_agent_extension=None): self._headers_multipart_form_data = { "User-Agent": self._generate_useragent(user_agent_extension), } + self.base_api_url = api_instance_url or SCALE_API_BASE_URL_V1 @staticmethod def _http_request( @@ -101,7 +102,7 @@ def _api_request( ): """Generic HTTP request method with error handling.""" - url = f"{SCALE_ENDPOINT}/{endpoint}" + url = f"{self.base_api_url}/{endpoint}" res = self._http_request(method, url, headers, auth, params, body, files, data)