Add counties

build.py

@@ -11,7 +11,7 @@ def dict_factory(cursor, row):
     return dict((col[0], row[idx]) for idx, col in enumerate(cursor.description))
 
 
-def pickle_data():
+def pickle_state_data():
 
     dbpath = os.path.abspath(os.path.join(PWD, 'data.db'))
 
@@ -38,8 +38,32 @@ def pickle_data():
         pickle.dump(states, pkl_file)
 
 
+def pickle_county_data():
+
+    dbpath = os.path.abspath(os.path.join(PWD, 'data.db'))
+
+    conn = sqlite3.connect(dbpath)
+    conn.row_factory = dict_factory
+
+    c = conn.cursor()
+    c.execute("""SELECT * FROM counties ORDER BY state_fips, name""")
+
+    counties = []
+
+    for row in c:
+        row['name_metaphone'] = jellyfish.metaphone(row['short_name'])
+        row['full_fips'] = row['state_fips'] + row['fips']
+        counties.append(row)
+
+    pkl_path = os.path.abspath(os.path.join(PWD, 'us', 'counties.pkl'))
+
+    with open(pkl_path, 'wb') as pkl_file:
+        pickle.dump(counties, pkl_file)
+
+
 def build():
-    pickle_data()
+    pickle_state_data()
+    pickle_county_data()
 
 
 if __name__ == '__main__':

county_schema.sql

@@ -0,0 +1,8 @@
+CREATE TABLE counties (
+    fips TEXT,
+    state_fips TEXT,
+    name TEXT,
+    short_name TEXT,
+    lsad TEXT,
+    UNIQUE (state_fips, fips)
+);

load_county_data.py

@@ -0,0 +1,39 @@
+from __future__ import print_function
+import requests
+import sqlite3
+import fiona
+from os import remove
+
+TEMP_DATA_FILE = u'county_shapefile.zip'
+SHP_URL = u'http://www2.census.gov/geo/tiger/TIGER2010/COUNTY/2010/tl_2010_us_county10.zip'
+SHP_FILENAME = u'tl_2010_us_county10.shp'
+
+def download_county_data():
+    r = requests.get(SHP_URL)
+    with open(TEMP_DATA_FILE, 'wb') as f:
+        for chunk in r.iter_content(chunk_size=128):
+            f.write(chunk)
+
+
+def feature_data(f):
+    p = 'properties'
+    return (f[p][u'COUNTYFP10'],
+            f[p][u'STATEFP10'],
+            f[p][u'NAMELSAD10'],
+            f[p][u'NAME10'],
+            f[p][u'LSAD10'])
+
+
+def load_county_data():
+    with fiona.open('/{}'.format(SHP_FILENAME), vfs='zip://{}'.format(TEMP_DATA_FILE)) as c:
+        data = [feature_data(feature) for feature in c]
+
+    with sqlite3.connect('data.db') as con:
+        q = "INSERT INTO counties(fips, state_fips, name, short_name, lsad) VALUES(?, ?, ?, ?, ?);"
+        con.executemany(q, data)
+
+
+if __name__ == '__main__':
+        download_county_data()
+        load_county_data()
+        remove(TEMP_DATA_FILE)

tests.py

@@ -98,6 +98,31 @@ def test_continental(self):
         # Lower 48 + DC + Alaska
         self.assertEqual(len(us.STATES_CONTINENTAL), 50)
 
+    def test_counties(self):
+        # 3221 counties
+        self.assertEqual(len(us.COUNTIES), 3221)
+
+
+class CountyLookupTestCase(unittest.TestCase):
+
+    def test_fips(self):
+        self.assertEqual(us.counties.lookup('27053')[0].name, "Hennepin County")
+
+    def test_name(self):
+        self.assertEqual(len(us.counties.lookup('Washington')), 31)
+        self.assertEqual(len(us.counties.lookup('Hennepin')), 1)
+
+    def test_fullfips(self):
+        self.assertEqual(us.counties.lookup('Hennepin')[0].full_fips, u'27053')
+
+class CountyStateLinkTestCase(unittest.TestCase):
+
+    def test_state_counties(self):
+        self.assertEqual(len(us.states.MN.counties), 87)
+
+    def test_county_state(self):
+        self.assertEqual(us.counties.lookup('27053')[0].state, us.states.MN)
+
 
 if __name__ == '__main__':
     unittest.main()

us/__init__.py

@@ -3,7 +3,17 @@
     STATES, STATES_CONTIGUOUS, STATES_CONTINENTAL,
     TERRITORIES, STATES_AND_TERRITORIES, OBSOLETE
 )
+from . import counties
+from .counties import COUNTIES
 from .unitedstatesofamerica import *
 
 __appname__ = 'us'
 __version__ = '0.9.1'
+
+def __link_states_and_counties():
+    for c in COUNTIES:
+        state = states.lookup(c.state_fips, field='fips')
+        c.state = state
+        state.counties.append(c)
+
+__link_states_and_counties()