feat: pycountry ISO codes

ecologits/electricity_mix_repository.py

@@ -1,149 +1,9 @@
 import os
 from csv import DictReader
 from dataclasses import dataclass
-from enum import Enum
 from typing import Optional
 
 
-class Zones(Enum):
-    world = "WOR"
-    europe = "EEE"
-    zimbabwe = "ZWE"
-    zambia = "ZMB"
-    south_africa = "ZAF"
-    yemen = "YEM"
-    vietnam = "VNM"
-    venezuela = "VEN"
-    uzbekistan = "UZB"
-    uruguay = "URY"
-    united_states = "USA"
-    ukraine = "UKR"
-    tanzania = "TZA"
-    taiwan = "TWN"
-    trinidad_and_tobago = "TTO"
-    turkey = "TUR"
-    tunisia = "TUN"
-    turkmenistan = "TKM"
-    tajikistan = "TJK"
-    thailand = "THA"
-    togo = "TGO"
-    syrian_arab_republic = "SYR"
-    el_salvador = "SLV"
-    senegal = "SEN"
-    slovak_republic = "SVK"
-    slovenia = "SVN"
-    singapore = "SGP"
-    sweden = "SWE"
-    sudan = "SDN"
-    saudi_arabia = "SAU"
-    russian_federation = "RUS"
-    serbia_and_montenegro = "SCG"
-    romania = "ROU"
-    qatar = "QAT"
-    paraguay = "PRY"
-    portugal = "PRT"
-    poland = "POL"
-    pakistan = "PAK"
-    philippines = "PHL"
-    peru = "PER"
-    panama = "PAN"
-    oman = "OMN"
-    new_zealand = "NZL"
-    nepal = "NPL"
-    norway = "NOR"
-    netherlands = "NLD"
-    nicaragua = "NIC"
-    nigeria = "NGA"
-    namibia = "NAM"
-    mozambique = "MOZ"
-    malaysia = "MYS"
-    mexico = "MEX"
-    malta = "MLT"
-    mongolia = "MNG"
-    myanmar = "MMR"
-    north_macedonia = "MKD"
-    moldova = "MDA"
-    morocco = "MAR"
-    libya = "LBY"
-    latvia = "LVA"
-    luxembourg = "LUX"
-    lithuania = "LTU"
-    sri_lanka = "LKA"
-    lebanon = "LBN"
-    kazakhstan = "KAZ"
-    kuwait = "KWT"
-    south_korea = "KOR"
-    north_korea = "PRK"
-    cambodia = "KHM"
-    kyrgyz_republic = "KGZ"
-    kenya = "KEN"
-    japan = "JPN"
-    jordan = "JOR"
-    jamaica = "JAM"
-    italy = "ITA"
-    iceland = "ISL"
-    iran = "IRN"
-    iraq = "IRQ"
-    india = "IND"
-    israel = "ISR"
-    ireland = "IRL"
-    indonesia = "IDN"
-    hungary = "HUN"
-    haiti = "HTI"
-    croatia = "HRV"
-    honduras = "HND"
-    guatemala = "GTM"
-    greece = "GRC"
-    ghana = "GHA"
-    georgia = "GEO"
-    united_kingdom = "GBR"
-    gabon = "GAB"
-    france = "FRA"
-    finland = "FIN"
-    ethiopia = "ETH"
-    spain = "ESP"
-    eritrea = "ERI"
-    egypt = "EGY"
-    estonia = "EST"
-    ecuador = "ECU"
-    algeria = "DZA"
-    dominican_republic = "DOM"
-    denmark = "DNK"
-    germany = "DEU"
-    czech_republic = "CZE"
-    cyprus = "CYP"
-    cuba = "CUB"
-    costa_rica = "CRI"
-    colombia = "COL"
-    china = "CHN"
-    cameroon = "CMR"
-    chile = "CHL"
-    cote_d_ivoire = "CIV"
-    switzerland = "CHE"
-    congo = "COG"
-    democratic_republic_of_the_congo = "COD"
-    canada = "CAN"
-    belarus = "BLR"
-    botswana = "BWA"
-    brazil = "BRA"
-    bolivia = "BOL"
-    brunei = "BRN"
-    benin = "BEN"
-    bahrain = "BHR"
-    bulgaria = "BGR"
-    belgium = "BEL"
-    bangladesh = "BGD"
-    bosnia_and_herzegovina = "BIH"
-    azerbaijan = "AZE"
-    australia = "AUS"
-    austria = "AUT"
-    argentina = "ARG"
-    angola = "AGO"
-    armenia = "ARM"
-    albania = "ALB"
-    united_arab_emirates = "ARE"
-
-
 @dataclass
 class ElectricityMix:
     zone: str
@@ -174,7 +34,7 @@ def from_csv(cls, filepath: Optional[str] = None) -> "ElectricityMixRepository":
             for row in csv:
                 electricity_mixes.append(
                     ElectricityMix(
-                        zone=Zones(row["name"]).name,
+                        zone=row["name"],
                         adpe=float(row["adpe"]),
                         pe=float(row["pe"]),
                         gwp=float(row["gwp"]),

ecologits/tracers/utils.py

@@ -1,5 +1,7 @@
 from typing import Optional
 
+from pycountry import countries
+
 from ecologits.electricity_mix_repository import electricity_mixes
 from ecologits.impacts.llm import compute_llm_impacts
 from ecologits.impacts.modeling import Impacts, Range
@@ -15,7 +17,7 @@ def llm_impacts(
     model_name: str,
     output_token_count: int,
     request_latency: float,
-    electricity_mix_zone: Optional[str] = "world",
+    electricity_mix_zone: Optional[str] = "WOR",
 ) -> Optional[Impacts]:
     """
     High-level function to compute the impacts of an LLM generation request.
@@ -25,11 +27,16 @@ def llm_impacts(
         model_name: Name of the LLM used.
         output_token_count: Number of generated tokens.
         request_latency: Measured request latency in seconds.
-        electricity_mix_zone: Electricity mix zone (world electricity mix by default).
+        electricity_mix_zone: ISO 3166-1 alpha-3 code of the electricity mix zone (world electricity mix by default).
 
     Returns:
         The impacts of an LLM generation request.
     """
+    if countries.get(alpha_3=electricity_mix_zone) is None:
+        # TODO: Replace with proper logging
+        print(f"`{electricity_mix_zone}` is currently not a valid ISO 3166-1 alpha-3 code")
+        return None
+
     model = models.find_model(provider=provider, model_name=model_name)
     if model is None:
         # TODO: Replace with proper logging