changed it to build dataset with a few more variables

[babynames-cdsw] / ssadata.py

diff --git a/ssadata.py b/ssadata.py
index 96055a3d77ca25e5d58ab6ad7f26f651e0f7b5e1..35fabe9435cc9ed292a26ddeb2af87de350cad70 100644 (file)
--- a/ssadata.py
+++ b/ssadata.py
@@ -1,13 +1,34 @@
-NAMES_LIST = "yob2013.txt"
+import glob
+import re
 
-boys = {}
-girls = {}
+def import_yob_file(filename):
+    boys = {}
+    girls = {}
 
-for line in open(NAMES_LIST, 'r').readlines():
-    name, gender, count = line.strip().split(",")
-    count = int(count)
+    with open(filename, "r") as f:
+        for line in f.readlines():
+            name, gender, count = line.strip().split(",")
+            count = int(count)
+            if gender == "F":
+                girls[name.lower()] = count
+            elif gender == "M":
+                boys[name.lower()] = count
+                
+    return((boys, girls))
 
-    if gender == "F":
-        girls[name.lower()] = count
-    elif gender == "M":
-        boys[name.lower()] = count
+years = {}
+for filename in glob.glob('yob*.txt'):
+    year = re.match(r'yob(\d{4})\.txt$', filename).group(1)
+    tmp_boys, tmp_girls = import_yob_file(filename)
+    years[year] = {'girls' : tmp_girls,
+                   'boys' : tmp_boys}
+
+## resort by year
+years = dict(sorted(years.items()))
+
+girls = years["2021"]["girls"]
+boys = years["2021"]["boys"]
+
+
+    
+