✨ Add excluding countries to harmonizer (#3483)

* ✨ Add excluding countries to harmonizer * wip * wip
owid · Nov 4, 2024 · 9bd5ecf · 9bd5ecf
1 parent 78f2fcd
commit 9bd5ecf
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 4 deletions.
diff --git a/apps/owidbot/cli.py b/apps/owidbot/cli.py
@@ -60,7 +60,8 @@ def cli(
     repo = gh_utils.get_repo(repo_name)
     pr = gh_utils.get_pr(repo, branch)
     if pr is None:
-        raise AssertionError(f"No open PR found for branch {branch}")
+        log.warning(f"No open PR found for branch {branch}")
+        return
 
     # recalculate services
     services_body = {}

diff --git a/etl/harmonize.py b/etl/harmonize.py
@@ -97,6 +97,9 @@ def harmonize(
     # Export
     harmonizer.export_mapping()
 
+    # Excluded countries
+    harmonizer.export_excluded_countries()
+
 
 def harmonize_ipython(
     tb: Table,
@@ -264,6 +267,7 @@ def __init__(
         # Mapping
         self._mapping = None
         self.countries_mapped_automatic = None
+        self.excluded = []
 
     def _get_geo(self, tb, colname, indicator):
         """Get set of country names to map."""
@@ -359,7 +363,8 @@ def run_interactive_terminal(
 
         questionary.print("Beginning interactive harmonization...")
         questionary.print("  Select [skip] to skip a country/region mapping")
-        questionary.print("  Select [custom] to enter a custom name\n")
+        questionary.print("  Select [custom] to enter a custom name")
+        questionary.print("  Select [exclude] to exclude a country\n")
 
         instruction = "(Use shortcuts or arrow keys)"
         # start interactive session
@@ -372,7 +377,7 @@ def run_interactive_terminal(
                 # show suggestions
                 name = questionary.select(
                     f"[{i}/{len(self.ambiguous)}] {region}:",
-                    choices=suggestions + ["[custom]", "[skip]"],
+                    choices=suggestions + ["[custom]", "[skip]", "[exclude]"],
                     use_shortcuts=True,
                     style=SHELL_FORM_STYLE,
                     instruction=instruction,
@@ -396,6 +401,9 @@ def run_interactive_terminal(
                 elif name == "[skip]":
                     n_skipped += 1
                     continue
+                elif name == "[exclude]":
+                    self.excluded.append(region)
+                    continue
 
                 self.mapping[region] = name
         except KeyboardInterrupt:
@@ -568,6 +576,14 @@ def export_mapping(self):
         with open(self.output_file, "w") as ostream:
             json.dump(self.mapping, ostream, indent=2)
 
+    def export_excluded_countries(self):
+        if self.output_file is None:
+            raise ValueError("`output_file` not provided")
+        assert ".countries." in str(self.output_file), "Output file is not in **/*.countries.json format"
+        excluded_countries_path = str(self.output_file).replace(".countries.", ".excluded_countries.")
+        with open(excluded_countries_path, "w") as ostream:
+            json.dump(self.excluded, ostream, indent=2)
+
 
 if __name__ == "__main__":
     harmonize()
diff --git a/tests/apps/wizard/app_pages/test_apps.py b/tests/apps/wizard/app_pages/test_apps.py
@@ -8,7 +8,7 @@
 from apps.wizard.utils import WIZARD_DIR
 from etl import config
 
-DEFAULT_TIMEOUT = 20
+DEFAULT_TIMEOUT = 30
 
 
 @contextmanager