<commit_before>import os
import csv

from ott.utils import file_utils
from ott.utils.cache_base import CacheBase


class Fix(CacheBase):
    """
    cmd line utility that opens a GTFS file, and can rename agency in routes/agency .txt files
    why this exists, I forget :-(  that said, I can think of a handful of times I had to manually fix a gtfs file, so...
    older comments below indicate this was created to match TriMet's gtfs-rt PSC alerts with the gtfs.zip data
    """
    gtfs_name = None
    gtfs_path = None

    def __init__(self, gtfs_name):
        super(Fix, self).__init__(section='gtfs')
        self.gtfs_name = gtfs_name
        self.gtfs_path = os.path.join(self.cache_dir, gtfs_name)

    def cp(self, to_gtfs_path=None):
        """ copy gtfs file to new file (good for testing, so you don't harm original gtfs.zip)
        """
        if to_gtfs_path is None:
            to_gtfs_path = self.gtfs_path + '.zip'
        file_utils.cp(self.gtfs_path, to_gtfs_path)
        self.gtfs_path = to_gtfs_path

    def rename_agency_in_routes_txt(self, regex_str, replace_str):
        file_utils.replace_strings_in_zipfile(self.gtfs_path, "routes.txt", regex_str, replace_str)

    def rename_agency_in_agency_txt(self, regex_str, replace_str):
        file_utils.replace_strings_in_zipfile(self.gtfs_path, "agency.txt", regex_str, replace_str)

    def remove_deadhead_stop_times(self, stop, file_name="stop_times.txt"):
        """
        will remove preceding stop_time in a trip given a target stop

        THIS:
          9683134,18:01:00,18:01:00,9654,1,Portland,1,0,0.0,1,,
          9683134,18:03:00,18:03:00,8169,2,Portland,0,1,267.5,1,,

        BECOMES:
          9683134,18:03:00,18:03:00,8169,1,Portland,0,0,0.0,1,,

        grep a stop from stop_times.txt (w/in a gtfs.zip feed file)
        writes a stop_times.txt.tmp file with
        got dict read/write ideas from:
          https://stackoverflow.com/questions/2982023/how-to-write-header-row-with-csv-dictwriter/2982117
        """

        # step 1: extract file from gtfs -- open temp file to write stuff
        in_file = file_utils.unzip_file(self.gtfs_path, file_name)

        # step 2: open this input csv
        with open(in_file, 'r') as csv_in:
            dr = csv.DictReader(csv_in)
            prev_row = None

            # step 3: open .csv output
            with open(in_file + ".tmp", 'w+') as csv_out:
                dw = csv.DictWriter(csv_out, fieldnames=dr.fieldnames)
                dw.writerow(dict(zip(dr.fieldnames, dr.fieldnames)))
                for row in dr:
                    # step 4: cull the preceding stop_time to our target stop, and move a few vars over
                    if prev_row and row.get('stop_id') == stop and row.get('trip_id') == prev_row.get('trip_id'):
                        row['stop_sequence'] = prev_row['stop_sequence']
                        row['shape_dist_traveled'] = prev_row['shape_dist_traveled']
                        row['pickup_type'] = "0"
                        row['drop_off_type'] = "0"
                        dw.writerow(row)
                        prev_row = None
                        continue

                    # step 5: write out normal (prev) row and move on to next row
                    if prev_row:
                        dw.writerow(prev_row)
                    prev_row = row

                # step 6: outside the read/write loop, write last row
                if prev_row:
                    dw.writerow(prev_row)

    @classmethod
    def get_args(cls):
        """ command-line arg parser and help util...
            examples:
               bin/gtfs_fix SAM.zip -a -r -f "^86" -t "SAM"
               bin/gtfs_fix SMART.zip -a -r -f "^108" -t "SMART"
        """
        import argparse
        parser = argparse.ArgumentParser(prog='gtfs-fix', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
        parser.add_argument('gtfs', help="Name of GTFS zip that is in the 'cache' (e.g., TRIMET.zip)")
        parser.add_argument('--regex',   '-f', help="string (regex) to find in files")
        parser.add_argument('--replace', '-t', help="string to replace found regex strings")
        parser.add_argument('--routes',  '-r', default=False, action='store_true', help='fix routes.txt')
        parser.add_argument('--agency',  '-a', default=False, action='store_true', help='fix agency.txt')
        parser.add_argument('--copy',    '-c', default=False, action='store_true', help="make a copy of the gtfs.zip (don't edit original")
        parser.add_argument('--stop',    '-s', default=False, help="test stop query")
        args = parser.parse_args()
        return args


def rename_trimet_agency():
    """ just rename routes.txt (good for OTP alerts) ... don't rename the agency.txt
        bin/gtfs_fix TRIMET.zip -r -f "(PSC|TRAM)" -t "TRIMET"
    """
    fix = Fix("TRIMET.zip")<commit_msg>fix(gtfs fix): rewrite stop_times.txt back to zip<commit_after>import os
import csv

from ott.utils import file_utils
from ott.utils.cache_base import CacheBase


class Fix(CacheBase):
    """
    cmd line utility that opens a GTFS file, and can rename agency in routes/agency .txt files
    why this exists, I forget :-(  that said, I can think of a handful of times I had to manually fix a gtfs file, so...
    older comments below indicate this was created to match TriMet's gtfs-rt PSC alerts with the gtfs.zip data
    """
    gtfs_name = None
    gtfs_path = None

    def __init__(self, gtfs_name):
        super(Fix, self).__init__(section='gtfs')
        self.gtfs_name = gtfs_name
        self.gtfs_path = os.path.join(self.cache_dir, gtfs_name)

    def cp(self, to_gtfs_path=None):
        """ copy gtfs file to new file (good for testing, so you don't harm original gtfs.zip)
        """
        if to_gtfs_path is None:
            to_gtfs_path = self.gtfs_path + '.zip'
        file_utils.cp(self.gtfs_path, to_gtfs_path)
        self.gtfs_path = to_gtfs_path

    def rename_agency_in_routes_txt(self, regex_str, replace_str):
        file_utils.replace_strings_in_zipfile(self.gtfs_path, "routes.txt", regex_str, replace_str)

    def rename_agency_in_agency_txt(self, regex_str, replace_str):
        file_utils.replace_strings_in_zipfile(self.gtfs_path, "agency.txt", regex_str, replace_str)

    def remove_deadhead_stop_times(self, stop, file_name="stop_times.txt", repack=True):
        """
        will remove preceding stop_time in a trip given a target stop

        THIS:
          9683134,18:01:00,18:01:00,9654,1,Portland,1,0,0.0,1,,
          9683134,18:03:00,18:03:00,8169,2,Portland,0,1,267.5,1,,

        BECOMES:
          9683134,18:03:00,18:03:00,8169,1,Portland,0,0,0.0,1,,

        grep a stop from stop_times.txt (w/in a gtfs.zip feed file)
        writes a stop_times.txt.tmp file with
        got dict read/write ideas from:
          https://stackoverflow.com/questions/2982023/how-to-write-header-row-with-csv-dictwriter/2982117
        """

        # step 1: extract file from gtfs -- open temp file to write stuff
        in_file = file_utils.unzip_file(self.gtfs_path, file_name)
        out_file = in_file + ".tmp"

        # step 2: open this input csv
        with open(in_file, 'r') as csv_in:
            dr = csv.DictReader(csv_in)
            prev_row = None

            # step 3: open .csv output
            with open(out_file, 'w+') as csv_out:
                dw = csv.DictWriter(csv_out, fieldnames=dr.fieldnames)
                dw.writerow(dict(zip(dr.fieldnames, dr.fieldnames)))
                for row in dr:
                    # step 4: cull the preceding stop_time to our target stop, and move a few vars over
                    if prev_row and row.get('stop_id') == stop and row.get('trip_id') == prev_row.get('trip_id'):
                        row['stop_sequence'] = prev_row['stop_sequence']
                        row['shape_dist_traveled'] = prev_row['shape_dist_traveled']
                        row['pickup_type'] = "0"
                        row['drop_off_type'] = "0"
                        dw.writerow(row)
                        prev_row = None
                        continue

                    # step 5: write out normal (prev) row and move on to next row
                    if prev_row:
                        dw.writerow(prev_row)
                    prev_row = row

                # step 6: outside the read/write loop, write last row
                if prev_row:
                    dw.writerow(prev_row)

        if repack:
            file_utils.replace_file_in_zipfile(self.gtfs_path, out_file, file_name)

    @classmethod
    def get_args(cls):
        """ command-line arg parser and help util...
            examples:
               bin/gtfs_fix SAM.zip -a -r -f "^86" -t "SAM"
               bin/gtfs_fix SMART.zip -a -r -f "^108" -t "SMART"
        """
        import argparse
        parser = argparse.ArgumentParser(prog='gtfs-fix', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
        parser.add_argument('gtfs', help="Name of GTFS zip that is in the 'cache' (e.g., TRIMET.zip)")
        parser.add_argument('--regex',   '-f', help="string (regex) to find in files")
        parser.add_argument('--replace', '-t', help="string to replace found regex strings")
        parser.add_argument('--routes',  '-r', default=False, action='store_true', help='fix routes.txt')
        parser.add_argument('--agency',  '-a', default=False, action='store_true', help='fix agency.txt')
        parser.add_argument('--copy',    '-c', default=False, action='store_true', help="make a copy of the gtfs.zip (don't edit original")
        parser.add_argument('--stop',    '-s', default=False, help="test stop query")
        args = parser.parse_args()
        return args


def rename_trimet_agency():
    """ just rename routes.txt (good for OTP alerts) ... don't rename the agency.txt
        bin/gtfs_fix TRIMET.zip -r -f "(PSC|TRAM)" -t "TRIMET"
    """
    fix = Fix("TRIMET.zip")
{'fraction_non_alphanumeric': 0.11272399639603564, 'fraction_numerical': 0.023225548102913204, 'mean_word_length': 2.2818659658344282, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, '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': '6474389', 'n_tokens_mistral': 3392, 'n_tokens_neox': 3140, 'n_words': 819}