Remove input country from output. Accept lower case country codes.
This commit is contained in:
@@ -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))
|
||||||
|
|||||||
Reference in New Issue
Block a user