Remove input country from output. Accept lower case country codes.

This commit is contained in:
Kira
2025-12-30 15:34:20 +01:00
committed by Nyanraltotlapun
parent 0311c41422
commit 2176c244ea
+8 -4
View File
@@ -182,7 +182,7 @@ NEAR_COUNTRIES = {
"YE": ['OM', 'SA'], "YE": ['OM', 'SA'],
"ZA": ['BW', 'LS', 'MZ', 'NA', 'SZ', 'ZW'], "ZA": ['BW', 'LS', 'MZ', 'NA', 'SZ', 'ZW'],
"ZM": ['AO', 'BW', 'CD', 'MW', 'MZ', 'NA', 'TZ', 'ZW'], "ZM": ['AO', 'BW', 'CD', 'MW', 'MZ', 'NA', 'TZ', 'ZW'],
"ZW": ['BW', 'MZ', 'ZA', 'ZM'], "ZW": ['BW', 'MZ', 'ZA', 'ZM']
} }
@@ -208,18 +208,22 @@ def main():
args = parser.parse_args() args = parser.parse_args()
search_country = args.country # Strip whitespaces and uppercase input
search_country = args.country.strip().upper()
recursion_depth = args.recursion recursion_depth = args.recursion
if recursion_depth > 12: if recursion_depth > 12:
sys.exit("Recursion should be less than or equal 12! You doing something stupid here.") sys.exit("Recursion should be less than or equal 12! You doing something stupid here.")
result = set(NEAR_COUNTRIES.get(search_country,[])) result = set(NEAR_COUNTRIES.get(search_country,[]))
for i in range(1, recursion_depth, 1): for i in range(1, recursion_depth):
n_c = list(result) n_c = list(result)
for country in n_c: for country in n_c:
result.update(NEAR_COUNTRIES.get(country,[])) result.update(NEAR_COUNTRIES.get(country,[]))
# Remove search country from result
result.discard(search_country)
list_result = list(result) list_result = list(result)
list_result.sort() list_result.sort()
sys.stdout.write(' '.join(list_result)) sys.stdout.write(' '.join(list_result))