Election Day: Add party results to archive.

TYPE: Feature
LINK: OGC-877

src/onegov/election_day/utils/archive_generator.py

@@ -10,8 +10,9 @@
 
 from onegov.core.csv import convert_list_of_dicts_to_csv
 from onegov.core.utils import module_path
-from onegov.ballot import Vote, Election, ElectionCompound
+from onegov.ballot import Vote, Election, ElectionCompound, ProporzElection
 from onegov.election_day.formats import export_internal
+from onegov.election_day.formats import export_parties_internal
 
 
 from typing import Any
@@ -82,12 +83,7 @@ def generate_csv(self) -> None:
                 year_dir = f'{entity_name}/{year}'
                 self.temp_fs.makedirs(year_dir, recreate=True)
                 for item in yearly_package:
-                    # item may be of type Vote, Election or ElectionCompound
-                    filename = item.id[: self.MAX_FILENAME_LENGTH] + '.csv'
-                    combined_path = path.combine(year_dir, filename)
-                    with self.temp_fs.open(combined_path, 'w') as f:
-                        rows = export_internal(item, sorted(self.app.locales))
-                        f.write(convert_list_of_dicts_to_csv(rows))
+                    self.export_item(item, year_dir)
 
         # Additionally, create 'flat csv' containing all votes in a single file
         if votes:
@@ -223,6 +219,31 @@ def include_docs(self) -> None:
                 dst_path=match.path,
             )
 
+    def export_item(self, item: 'EntityT', dir: str) -> None:
+        locales = sorted(self.app.locales)
+        default_locale = self.app.default_locale
+
+        # results
+        filename = item.id[:self.MAX_FILENAME_LENGTH] + '.csv'
+        combined_path = path.combine(dir, filename)
+        rows = export_internal(item, locales)
+        with self.temp_fs.open(combined_path, 'w') as f:
+            f.write(convert_list_of_dicts_to_csv(rows))
+
+        # party results
+        if getattr(item, 'has_party_results', False):
+            assert isinstance(item, (ProporzElection, ElectionCompound))
+            filename = item.id[:self.MAX_FILENAME_LENGTH + 8] + '-parties.csv'
+            combined_path = path.combine(dir, filename)
+            rows = export_parties_internal(
+                item,
+                locales,
+                # FIXME: Should we assert that the default_locale is set?
+                default_locale=default_locale,  # type:ignore[arg-type]
+            )
+            with self.temp_fs.open(combined_path, 'w') as f:
+                f.write(convert_list_of_dicts_to_csv(rows))
+
     def generate_archive(self) -> str | None:
         self.generate_csv()
         self.include_docs()

tests/onegov/election_day/utils/test_archive_generator.py

@@ -224,7 +224,7 @@ def test_long_filenames_are_truncated(election_day_app_zg):
 
 
 def test_election_generation(election_day_app_zg, import_test_datasets):
-    import_test_datasets(
+    election, errors = import_test_datasets(
         'internal',
         'election',
         'zg',
@@ -235,6 +235,17 @@ def test_election_generation(election_day_app_zg, import_test_datasets):
         dataset_name='nationalratswahlen-2015',
         app_session=election_day_app_zg.session()
     )
+    assert not errors
+    errors = import_test_datasets(
+        'internal',
+        'parties',
+        'zg',
+        'canton',
+        'proporz',
+        election=election,
+        dataset_name='nationalratswahlen-2015-parteien',
+    )
+    assert not errors
 
     archive_generator = ArchiveGenerator(election_day_app_zg)
     zip_path = archive_generator.generate_archive()
@@ -244,13 +255,14 @@ def test_election_generation(election_day_app_zg, import_test_datasets):
             assert "elections" in top_level_dir
 
             elections = zip_fs.opendir("elections")
-            elections_by_year = [year for year in elections.listdir(".")]
-
-            assert {"2015"} == set(elections_by_year)
-
-            csv = [csv for csv in zip_fs.scandir("elections/2015",
-                                                 namespaces=["basic"])]
-            first_file = csv[0]
-
-            assert "proporz_internal_nationalratswahlen-2015.csv" ==\
-                   first_file.name
+            years = {year for year in elections.listdir(".")}
+            assert years == {"2015"}
+
+            files = {
+                csv.name for csv
+                in zip_fs.scandir("elections/2015", namespaces=["basic"])
+            }
+            assert files == {
+                'proporz_internal_nationalratswahlen-2015.csv',
+                'proporz_internal_nationalratswahlen-2015-parties.csv'
+            }