diff --git a/codecarbon/external/geography.py b/codecarbon/external/geography.py index c300d2f92..075824959 100644 --- a/codecarbon/external/geography.py +++ b/codecarbon/external/geography.py @@ -3,10 +3,10 @@ """ import re -import urllib.parse from dataclasses import dataclass from typing import Callable, Dict, Optional +import pycountry import requests from codecarbon.core.cloud import get_env_cloud_details @@ -93,10 +93,14 @@ def from_geo_js(cls, url: str) -> "GeoMetadata": try: response: Dict = requests.get(url, timeout=0.5).json() + region = response.get("region", "").lower() + if not region: + raise ValueError("Region is empty") + return cls( country_iso_code=response["country_code3"].upper(), country_name=response["country"], - region=response.get("region", "").lower(), + region=region, latitude=float(response.get("latitude")), longitude=float(response.get("longitude")), country_2letter_iso_code=response.get("country_code"), @@ -107,32 +111,36 @@ def from_geo_js(cls, url: str) -> "GeoMetadata": f"Unable to access geographical location through primary API. Will resort to using the backup API - Exception : {e} - url={url}" ) - geo_url_backup = "https://ip-api.com/json/" + geo_url_backup = "https://ipinfo.io/json" try: geo_response: Dict = requests.get(geo_url_backup, timeout=0.5).json() - country_name = geo_response["country"] - # The previous request does not return the three-letter country code - country_code_3_url = f"https://api.first.org/data/v1/countries?q={urllib.parse.quote_plus(country_name)}&scope=iso" - country_code_response: Dict = requests.get( - country_code_3_url, timeout=0.5 - ).json() + # extract latitude and longitude from loc (e.g., "loc": "37.4056,-122.0775") + loc = geo_response.get("loc", "").split(",") + latitude = float(loc[0]) if len(loc) == 2 else 0.0 + longitude = float(loc[1]) if len(loc) == 2 else 0.0 + + # Retrieve the 3-letter ISO code using pycountry + country_2letter_iso_code = geo_response.get("country") + country = pycountry.countries.get(alpha_2=country_2letter_iso_code) + + # Some countries might not be found or mapped perfectly + country_iso_code = country.alpha_3 if country else "" + country_name = country.name if country else "" return cls( - country_iso_code=next( - iter(country_code_response["data"].keys()) - ).upper(), + country_iso_code=country_iso_code.upper(), country_name=country_name, - region=geo_response.get("regionName", "").lower(), - latitude=float(geo_response.get("lat")), - longitude=float(geo_response.get("lon")), - country_2letter_iso_code=geo_response.get("countryCode"), + region=geo_response.get("region", "").lower(), + latitude=latitude, + longitude=longitude, + country_2letter_iso_code=country_2letter_iso_code, ) except Exception as e: # If both API calls fail, default to Canada logger.warning( - f"Unable to access geographical location. Using 'Canada' as the default value - Exception : {e} - url={url}" + f"Unable to access geographical location through fallback API. Using 'Canada' as the default value - Exception : {e} - url={geo_url_backup}" ) return cls( diff --git a/pyproject.toml b/pyproject.toml index 518acb7ed..028587aed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ dependencies = [ "questionary", "rich", "typer", + "pycountry", ] [tool.setuptools.dynamic] diff --git a/tests/test_geography.py b/tests/test_geography.py index b92a79f0d..8f95f7f43 100644 --- a/tests/test_geography.py +++ b/tests/test_geography.py @@ -9,7 +9,6 @@ CLOUD_METADATA_AZURE, CLOUD_METADATA_GCP, CLOUD_METADATA_GCP_EMPTY, - COUNTRY_METADATA_USA, GEO_METADATA_CANADA, GEO_METADATA_USA, GEO_METADATA_USA_BACKUP, @@ -91,14 +90,27 @@ def test_geo_metadata_USA_backup(self): ) responses.add( responses.GET, - "https://ip-api.com/json/", + "https://ipinfo.io/json", json=GEO_METADATA_USA_BACKUP, status=200, ) + geo = GeoMetadata.from_geo_js(self.geo_js_url) + self.assertEqual("USA", geo.country_iso_code) + self.assertEqual("United States", geo.country_name) + self.assertEqual("illinois", geo.region) + + @responses.activate + def test_geo_metadata_empty_region_fallback(self): + empty_region_response = GEO_METADATA_USA.copy() + empty_region_response["region"] = "" + + responses.add( + responses.GET, self.geo_js_url, json=empty_region_response, status=200 + ) responses.add( responses.GET, - "https://api.first.org/data/v1/countries?q=United%20States&scope=iso", - json=COUNTRY_METADATA_USA, + "https://ipinfo.io/json", + json=GEO_METADATA_USA_BACKUP, status=200, ) geo = GeoMetadata.from_geo_js(self.geo_js_url) diff --git a/tests/testdata.py b/tests/testdata.py index b4b70f94b..c70dd10eb 100644 --- a/tests/testdata.py +++ b/tests/testdata.py @@ -19,38 +19,14 @@ } GEO_METADATA_USA_BACKUP = { - "organization_name": "foobar", - "regionName": "Illinois", - "accuracy": 1, - "asn": 0, - "organization": "foobar", - "timezone": "America/Chicago", - "lon": "88", - "area_code": "0", "ip": "foobar", "city": "Chicago", - "country": "United States", - "countryCode": "US", - "lat": "0", -} - -COUNTRY_METADATA_USA = { - "status": "OK", - "status-code": 200, - "version": "1.0", - "access": "public", - "data": { - "USA": { - "id": "USA", - "country": "United States of America (the)", - "region": "North America", - }, - "UMI": { - "id": "UMI", - "country": "United States Minor Outlying Islands (the)", - "region": "Oceania", - }, - }, + "region": "Illinois", + "country": "US", + "loc": "0,88", + "org": "foobar", + "postal": "60601", + "timezone": "America/Chicago", } GEO_METADATA_CANADA = { diff --git a/uv.lock b/uv.lock index 8161de636..3b4a81200 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.10" resolution-markers = [ "python_full_version < '3.11'", @@ -332,6 +332,7 @@ dependencies = [ { name = "prometheus-client" }, { name = "psutil" }, { name = "py-cpuinfo" }, + { name = "pycountry" }, { name = "pydantic" }, { name = "questionary" }, { name = "rapidfuzz" }, @@ -391,6 +392,7 @@ requires-dist = [ { name = "prometheus-client" }, { name = "psutil", specifier = ">=6.0.0" }, { name = "py-cpuinfo" }, + { name = "pycountry" }, { name = "pydantic" }, { name = "questionary" }, { name = "rapidfuzz" }, @@ -1613,6 +1615,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/a9/023730ba63db1e494a271cb018dcd361bd2c917ba7004c3e49d5daf795a2/py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5", size = 22335, upload-time = "2022-10-25T20:38:27.636Z" }, ] +[[package]] +name = "pycountry" +version = "26.2.16" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/de/1d/061b9e7a48b85cfd69f33c33d2ef784a531c359399ad764243399673c8f5/pycountry-26.2.16.tar.gz", hash = "sha256:5b6027d453fcd6060112b951dd010f01f168b51b4bf8a1f1fc8c95c8d94a0801", size = 7711342, upload-time = "2026-02-17T03:42:52.367Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/42/7703bd45b62fecd44cd7d3495423097e2f7d28bc2e99e7c1af68892ab157/pycountry-26.2.16-py3-none-any.whl", hash = "sha256:115c4baf7cceaa30f59a4694d79483c9167dbce7a9de4d3d571c5f3ea77c305a", size = 8044600, upload-time = "2026-02-17T03:42:49.777Z" }, +] + [[package]] name = "pycparser" version = "3.0"