<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import csv, json
from glob import glob
from os.path import basename, join

with open(join('us-data', 'codes.txt')) as f:
    rows = list(csv.DictReader(f, dialect='excel-tab'))
    codes = dict([(row['Postal Code'].lower(), row['State']) for row in rows])

with open(join('us-data', 'states.txt')) as f:
    rows = list(csv.DictReader(f, dialect='excel-tab'))
    states = dict([(row['Name'], row['State FIPS']) for row in rows])

with open(join('us-data', 'counties.txt')) as f:
    counties = dict()
    
    for row in csv.DictReader(f, dialect='excel-tab'):
        key = row['State FIPS'], row['Name']
        value = row['County FIPS'], row['Name']
        counties[key] = value
        
        # some key variations
        if row['Name'].endswith(' County'):
            counties[(row['State FIPS'], row['Name'][:-7])] = value
        if row['Name'].endswith(' Parish'):
            counties[(row['State FIPS'], row['Name'][:-7])] = value
        if row['Name'].endswith(' Municipality'):
            counties[(row['State FIPS'], row['Name'][:-13])] = value
    
    # more key variations
    for ((s, c), value) in list(counties.items()):
        if c.startswith('St. '):
            counties[(s, 'Saint '+c[4:])] = value
    
    for ((s, c), value) in list(counties.items()):
        counties[(s, c.lower())] = value
        counties[(s, c.replace('-', ' '))] = value
        counties[(s, c.replace('-', ' ').lower())] = value

for path in glob('sources/us-*.json'):
    with open(path) as f:
        data = f.read()
        info = json.loads(data)
    
    if 'county' not in info.get('coverage', {}):
        continue
    
    if 'US Census' in info['coverage']:
        continue
    
    assert '\n    "coverage": {\n        "' in data
    
    print path, '...'
    
    state_name = codes[info['coverage']['state']]
    state_fips = states[state_name]
    
    county = info['coverage']['county']
    
    # if type(county) is list or basename(path)[6:-5] != county.lower().replace(' ', '-'):
    #     continue
    
    if type(county) is list:
        county_names = [counties[(state_fips, c)] for c in county]
        print info['coverage'], state_fips, state_name, county_names
        continue

    if u'ñ' in county:
        county_fips, county_name = counties[(state_fips, county.replace(u'ñ', 'n'))]
    else:
        county_fips, county_name = counties[(state_fips, county)]
    
    geoid = state_fips + county_fips
    
    census_dict = dict(geoid=geoid, name=county_name, state=state_name)
    census_json = json.dumps(census_dict, sort_keys=True)
    
    new_data = data.replace('\n    "coverage": {\n        "', '\n    "coverage": {{\n        "US Census": {0},\n        "'.format(census_json))
    
    with open(path, 'w') as file:
        file.write(new_data)
<commit_msg>Tweaked prefix search, to shorten line length<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import csv, json
from glob import glob
from os.path import basename, join

with open(join('us-data', 'codes.txt')) as f:
    rows = list(csv.DictReader(f, dialect='excel-tab'))
    codes = dict([(row['Postal Code'].lower(), row['State']) for row in rows])

with open(join('us-data', 'states.txt')) as f:
    rows = list(csv.DictReader(f, dialect='excel-tab'))
    states = dict([(row['Name'], row['State FIPS']) for row in rows])

with open(join('us-data', 'counties.txt')) as f:
    counties = dict()
    
    for row in csv.DictReader(f, dialect='excel-tab'):
        key = row['State FIPS'], row['Name']
        value = row['County FIPS'], row['Name']
        counties[key] = value
        
        # some key variations
        if row['Name'].endswith(' County'):
            counties[(row['State FIPS'], row['Name'][:-7])] = value
        if row['Name'].endswith(' Parish'):
            counties[(row['State FIPS'], row['Name'][:-7])] = value
        if row['Name'].endswith(' Municipality'):
            counties[(row['State FIPS'], row['Name'][:-13])] = value
    
    # more key variations
    for ((s, c), value) in list(counties.items()):
        if c.startswith('St. '):
            counties[(s, 'Saint '+c[4:])] = value
    
    for ((s, c), value) in list(counties.items()):
        counties[(s, c.lower())] = value
        counties[(s, c.replace('-', ' '))] = value
        counties[(s, c.replace('-', ' ').lower())] = value

for path in glob('sources/us-*.json'):
    with open(path) as f:
        data = f.read()
        info = json.loads(data)
    
    if 'county' not in info.get('coverage', {}):
        continue
    
    if 'US Census' in info['coverage']:
        continue
    
    prefix = '\n    "coverage": {\n        '
    assert prefix + '"' in data
    
    print path, '...'
    
    state_name = codes[info['coverage']['state']]
    state_fips = states[state_name]
    
    county = info['coverage']['county']
    
    # if type(county) is list or basename(path)[6:-5] != county.lower().replace(' ', '-'):
    #     continue
    
    if type(county) is list:
        county_names = [counties[(state_fips, c)] for c in county]
        print info['coverage'], state_fips, state_name, county_names
        continue

    if u'ñ' in county:
        county_fips, county_name = counties[(state_fips, county.replace(u'ñ', 'n'))]
    else:
        county_fips, county_name = counties[(state_fips, county)]
    
    geoid = state_fips + county_fips
    
    census_dict = dict(geoid=geoid, name=county_name, state=state_name)
    census_json = json.dumps(census_dict, sort_keys=True)
    
    new_data = data.replace(prefix, '{0}"US Census": {1},\n        '.format(prefix, census_json))
    
    with open(path, 'w') as file:
        file.write(new_data)

{'fraction_non_alphanumeric': 0.18037919638197947, 'fraction_numerical': 0.0033049225952339538, 'mean_word_length': 2.655435473617292, 'pattern_counts': {'":': 6, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3343523', 'n_tokens_mistral': 1947, 'n_tokens_neox': 1933, 'n_words': 462}