From efea7334473330edb4cbb9ace01a868c9f6f4e19 Mon Sep 17 00:00:00 2001
From: Felix Blanke <felixblanke@uni-bonn.de>
Date: Sun, 27 Aug 2023 00:00:02 +0200
Subject: [PATCH] Add command for special tag

---
 download_digital.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/download_digital.py b/download_digital.py
index f11e2ac..4a563c5 100644
--- a/download_digital.py
+++ b/download_digital.py
@@ -75,6 +75,7 @@ def construct_dataframe(
     url: str = "https://beschaeftigtenbefragung.verdi.de/",
     tag: str = "bez_data_2",
     grouped: bool = False,
+    special_tag: str | None = None,
 ):
     r = requests.get(url)
     soup = BeautifulSoup(r.text, "html.parser")
@@ -97,10 +98,16 @@ def construct_dataframe(
 
     tot_col_data = []
     tot_col_index = []
-    for k, v in bez_data.items():
-        if "tot" in v:
-            tot_col_data.append(v["tot"])
-            tot_col_index.append(k)
+    if special_tag:
+        for k, v in bez_data.items():
+            if "sp" in v and special_tag in v["sp"]:
+                tot_col_data.append(v["sp"][special_tag])
+                tot_col_index.append(k)
+    else:
+        for k, v in bez_data.items():
+            if "tot" in v:
+                tot_col_data.append(v["tot"])
+                tot_col_index.append(k)
 
     data["Digitale Befragung"] = pd.Series(tot_col_data, index=tot_col_index)
     df = pd.DataFrame(data=data)
@@ -116,8 +123,9 @@ def main(
     tag: str = "bez_data_2",
     dry_run: bool = False,
     grouped: bool = False,
+    special_tag: str | None = None,
 ) -> None:
-    df = construct_dataframe(url=url, tag=tag, grouped=grouped)
+    df = construct_dataframe(url=url, tag=tag, grouped=grouped, special_tag=special_tag)
 
     if dry_run:
         print(df)