226 lines
8.3 KiB
Python
226 lines
8.3 KiB
Python
"""
|
|
Utility data models and classes for common tools and mappings.
|
|
"""
|
|
|
|
class CountryCodes:
|
|
"""
|
|
Centralized country code mapping for different services.
|
|
|
|
Maps ISO-2 country codes to service-specific country names.
|
|
Each service may have different requirements for country names.
|
|
"""
|
|
|
|
# Mapping: ISO-2 code -> (Tavily country name, Perplexity country name)
|
|
_COUNTRY_MAP = {
|
|
"AF": ("afghanistan", "Afghanistan"),
|
|
"AL": ("albania", "Albania"),
|
|
"DZ": ("algeria", "Algeria"),
|
|
"AD": ("andorra", "Andorra"),
|
|
"AO": ("angola", "Angola"),
|
|
"AR": ("argentina", "Argentina"),
|
|
"AM": ("armenia", "Armenia"),
|
|
"AU": ("australia", "Australia"),
|
|
"AT": ("austria", "Austria"),
|
|
"AZ": ("azerbaijan", "Azerbaijan"),
|
|
"BS": ("bahamas", "Bahamas"),
|
|
"BH": ("bahrain", "Bahrain"),
|
|
"BD": ("bangladesh", "Bangladesh"),
|
|
"BB": ("barbados", "Barbados"),
|
|
"BY": ("belarus", "Belarus"),
|
|
"BE": ("belgium", "Belgium"),
|
|
"BZ": ("belize", "Belize"),
|
|
"BJ": ("benin", "Benin"),
|
|
"BT": ("bhutan", "Bhutan"),
|
|
"BO": ("bolivia", "Bolivia"),
|
|
"BA": ("bosnia and herzegovina", "Bosnia and Herzegovina"),
|
|
"BW": ("botswana", "Botswana"),
|
|
"BR": ("brazil", "Brazil"),
|
|
"BN": ("brunei", "Brunei"),
|
|
"BG": ("bulgaria", "Bulgaria"),
|
|
"BF": ("burkina faso", "Burkina Faso"),
|
|
"BI": ("burundi", "Burundi"),
|
|
"KH": ("cambodia", "Cambodia"),
|
|
"CM": ("cameroon", "Cameroon"),
|
|
"CA": ("canada", "Canada"),
|
|
"CV": ("cape verde", "Cape Verde"),
|
|
"CF": ("central african republic", "Central African Republic"),
|
|
"TD": ("chad", "Chad"),
|
|
"CL": ("chile", "Chile"),
|
|
"CN": ("china", "China"),
|
|
"CO": ("colombia", "Colombia"),
|
|
"KM": ("comoros", "Comoros"),
|
|
"CG": ("congo", "Congo"),
|
|
"CR": ("costa rica", "Costa Rica"),
|
|
"HR": ("croatia", "Croatia"),
|
|
"CU": ("cuba", "Cuba"),
|
|
"CY": ("cyprus", "Cyprus"),
|
|
"CZ": ("czech republic", "Czech Republic"),
|
|
"DK": ("denmark", "Denmark"),
|
|
"DJ": ("djibouti", "Djibouti"),
|
|
"DO": ("dominican republic", "Dominican Republic"),
|
|
"EC": ("ecuador", "Ecuador"),
|
|
"EG": ("egypt", "Egypt"),
|
|
"SV": ("el salvador", "El Salvador"),
|
|
"GQ": ("equatorial guinea", "Equatorial Guinea"),
|
|
"ER": ("eritrea", "Eritrea"),
|
|
"EE": ("estonia", "Estonia"),
|
|
"ET": ("ethiopia", "Ethiopia"),
|
|
"FJ": ("fiji", "Fiji"),
|
|
"FI": ("finland", "Finland"),
|
|
"FR": ("france", "France"),
|
|
"GA": ("gabon", "Gabon"),
|
|
"GM": ("gambia", "Gambia"),
|
|
"GE": ("georgia", "Georgia"),
|
|
"DE": ("germany", "Germany"),
|
|
"GH": ("ghana", "Ghana"),
|
|
"GR": ("greece", "Greece"),
|
|
"GT": ("guatemala", "Guatemala"),
|
|
"GN": ("guinea", "Guinea"),
|
|
"HT": ("haiti", "Haiti"),
|
|
"HN": ("honduras", "Honduras"),
|
|
"HU": ("hungary", "Hungary"),
|
|
"IS": ("iceland", "Iceland"),
|
|
"IN": ("india", "India"),
|
|
"ID": ("indonesia", "Indonesia"),
|
|
"IR": ("iran", "Iran"),
|
|
"IQ": ("iraq", "Iraq"),
|
|
"IE": ("ireland", "Ireland"),
|
|
"IL": ("israel", "Israel"),
|
|
"IT": ("italy", "Italy"),
|
|
"JM": ("jamaica", "Jamaica"),
|
|
"JP": ("japan", "Japan"),
|
|
"JO": ("jordan", "Jordan"),
|
|
"KZ": ("kazakhstan", "Kazakhstan"),
|
|
"KE": ("kenya", "Kenya"),
|
|
"KW": ("kuwait", "Kuwait"),
|
|
"KG": ("kyrgyzstan", "Kyrgyzstan"),
|
|
"LV": ("latvia", "Latvia"),
|
|
"LB": ("lebanon", "Lebanon"),
|
|
"LS": ("lesotho", "Lesotho"),
|
|
"LR": ("liberia", "Liberia"),
|
|
"LY": ("libya", "Libya"),
|
|
"LI": ("liechtenstein", "Liechtenstein"),
|
|
"LT": ("lithuania", "Lithuania"),
|
|
"LU": ("luxembourg", "Luxembourg"),
|
|
"MG": ("madagascar", "Madagascar"),
|
|
"MW": ("malawi", "Malawi"),
|
|
"MY": ("malaysia", "Malaysia"),
|
|
"MV": ("maldives", "Maldives"),
|
|
"ML": ("mali", "Mali"),
|
|
"MT": ("malta", "Malta"),
|
|
"MR": ("mauritania", "Mauritania"),
|
|
"MU": ("mauritius", "Mauritius"),
|
|
"MX": ("mexico", "Mexico"),
|
|
"MD": ("moldova", "Moldova"),
|
|
"MC": ("monaco", "Monaco"),
|
|
"MN": ("mongolia", "Mongolia"),
|
|
"ME": ("montenegro", "Montenegro"),
|
|
"MA": ("morocco", "Morocco"),
|
|
"MZ": ("mozambique", "Mozambique"),
|
|
"MM": ("myanmar", "Myanmar"),
|
|
"NA": ("namibia", "Namibia"),
|
|
"NP": ("nepal", "Nepal"),
|
|
"NL": ("netherlands", "Netherlands"),
|
|
"NZ": ("new zealand", "New Zealand"),
|
|
"NI": ("nicaragua", "Nicaragua"),
|
|
"NE": ("niger", "Niger"),
|
|
"NG": ("nigeria", "Nigeria"),
|
|
"KP": ("north korea", "North Korea"),
|
|
"MK": ("north macedonia", "North Macedonia"),
|
|
"NO": ("norway", "Norway"),
|
|
"OM": ("oman", "Oman"),
|
|
"PK": ("pakistan", "Pakistan"),
|
|
"PA": ("panama", "Panama"),
|
|
"PG": ("papua new guinea", "Papua New Guinea"),
|
|
"PY": ("paraguay", "Paraguay"),
|
|
"PE": ("peru", "Peru"),
|
|
"PH": ("philippines", "Philippines"),
|
|
"PL": ("poland", "Poland"),
|
|
"PT": ("portugal", "Portugal"),
|
|
"QA": ("qatar", "Qatar"),
|
|
"RO": ("romania", "Romania"),
|
|
"RU": ("russia", "Russia"),
|
|
"RW": ("rwanda", "Rwanda"),
|
|
"SA": ("saudi arabia", "Saudi Arabia"),
|
|
"SN": ("senegal", "Senegal"),
|
|
"RS": ("serbia", "Serbia"),
|
|
"SG": ("singapore", "Singapore"),
|
|
"SK": ("slovakia", "Slovakia"),
|
|
"SI": ("slovenia", "Slovenia"),
|
|
"SO": ("somalia", "Somalia"),
|
|
"ZA": ("south africa", "South Africa"),
|
|
"KR": ("south korea", "South Korea"),
|
|
"SS": ("south sudan", "South Sudan"),
|
|
"ES": ("spain", "Spain"),
|
|
"LK": ("sri lanka", "Sri Lanka"),
|
|
"SD": ("sudan", "Sudan"),
|
|
"SE": ("sweden", "Sweden"),
|
|
"CH": ("switzerland", "Switzerland"),
|
|
"SY": ("syria", "Syria"),
|
|
"TW": ("taiwan", "Taiwan"),
|
|
"TJ": ("tajikistan", "Tajikistan"),
|
|
"TZ": ("tanzania", "Tanzania"),
|
|
"TH": ("thailand", "Thailand"),
|
|
"TG": ("togo", "Togo"),
|
|
"TT": ("trinidad and tobago", "Trinidad and Tobago"),
|
|
"TN": ("tunisia", "Tunisia"),
|
|
"TR": ("turkey", "Turkey"),
|
|
"TM": ("turkmenistan", "Turkmenistan"),
|
|
"UG": ("uganda", "Uganda"),
|
|
"UA": ("ukraine", "Ukraine"),
|
|
"AE": ("united arab emirates", "United Arab Emirates"),
|
|
"GB": ("united kingdom", "United Kingdom"),
|
|
"US": ("united states", "United States"),
|
|
"UY": ("uruguay", "Uruguay"),
|
|
"UZ": ("uzbekistan", "Uzbekistan"),
|
|
"VE": ("venezuela", "Venezuela"),
|
|
"VN": ("vietnam", "Vietnam"),
|
|
"YE": ("yemen", "Yemen"),
|
|
"ZM": ("zambia", "Zambia"),
|
|
"ZW": ("zimbabwe", "Zimbabwe"),
|
|
}
|
|
|
|
@classmethod
|
|
def getForTavily(cls, isoCode: str) -> str:
|
|
"""
|
|
Get Tavily-compatible country name from ISO-2 code.
|
|
|
|
Args:
|
|
isoCode: ISO-2 country code (e.g., "CH", "ch", "US", "us")
|
|
|
|
Returns:
|
|
Country name in lowercase as required by Tavily (e.g., "switzerland", "united states")
|
|
"""
|
|
# Convert to uppercase for lookup
|
|
isoCodeUpper = isoCode.upper() if isoCode else ""
|
|
mapping = cls._COUNTRY_MAP.get(isoCodeUpper)
|
|
return mapping[0] if mapping else isoCode
|
|
|
|
@classmethod
|
|
def getForPerplexity(cls, isoCode: str) -> str:
|
|
"""
|
|
Get Perplexity-compatible country name from ISO-2 code.
|
|
|
|
Args:
|
|
isoCode: ISO-2 country code (e.g., "CH", "US")
|
|
|
|
Returns:
|
|
Full country name as required by Perplexity (e.g., "Switzerland", "United States")
|
|
"""
|
|
isoCodeUpper = isoCode.upper()
|
|
mapping = cls._COUNTRY_MAP.get(isoCodeUpper)
|
|
return mapping[1] if mapping else isoCode
|
|
|
|
@classmethod
|
|
def isValid(cls, isoCode: str) -> bool:
|
|
"""
|
|
Check if ISO-2 code is valid.
|
|
|
|
Args:
|
|
isoCode: ISO-2 country code to check
|
|
|
|
Returns:
|
|
True if valid, False otherwise
|
|
"""
|
|
return isoCode.upper() in cls._COUNTRY_MAP
|
|
|