upgrade_translations.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-

# Copyright (C) 2019 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Upgrades GTFS from Google translations extension [1] to GTFS-Translations [2].

[1] http://developers.google.com/transit/gtfs/reference/gtfs-extensions#translationstxt
[2] http://bit.ly/gtfs-translations

Usage.

Upgrade translations of a feed unpacked to `my-feed` directory and store them
to `my-feed_updated`:

  $ upgrade_translations.py my-feed

Specify output directory name explicitly:

  $ upgrade_translations.py my-feed-old my-feed-new

Sample feed.

  feed_info.txt:
  feed_publisher_name,feed_publisher_url,feed_lang
  Narnia,http://en.wikipedia.org/wiki/Narnia,en

  stops.txt:
  stop_id,stop_name,stop_lat,stop_lon
  stop1,Palace,10,11

  trips.txt:
  route_id,service_id,trip_id,trip_headsign
  sledge,service1,trip1,To Palace

Translations in Google extension format.

  translations.txt:
  trans_id,lang,translation
  http://en.wikipedia.org/wiki/Narnia,en,http://en.wikipedia.org/wiki/Narnia
  http://en.wikipedia.org/wiki/Narnia,es,http://es.wikipedia.org/wiki/Narnia
  Palace,en,Palace
  Palace,es,Palacio
  To Palace,en,To Palace
  To Palace,es,Palacio

Translations in GTFS-Translations format.

  translations.txt:
  table_name,field_name,language,translation,record_id,record_sub_id,field_value
  feed_info,feed_publisher_url,es,http://es.wikipedia.org/wiki/Narnia,,,
  stops,stop_name,es,Palacio,stop1,,
  trips,trip_headsign,es,Palacio,,,To Palace
"""

import csv
import os
import os.path
import shutil
import sys

# GTFS-Translations defines record_id and record_sub_id used for referencing a
# row in a GTFS table that requires translation.
RECORD_ID_MAP = {
    "agency": ("agency_id", None),
    "stops": ("stop_id", None),
    "routes": ("route_id", None),
    "trips": ("trip_id", None),
    "stop_times": ("trip_id", "stop_sequence"),
    "feed_info": (None, None),
    "calendar": ("service_id", None),
    "calendar_dates": ("service_id", "date"),
    "fare_attributes": ("fare_id", None),
    "fare_rules": ("fare_id", "route_id"),
    "shapes": ("shape_id", "shape_pt_sequence"),
    "frequencies": ("trip_id", "start_time"),
    "transfers": ("from_stop_id", "to_stop_id"),
    "pathways": ("pathway_id", None),
    "levels": ("level_id", None),
}

# File translations.txt in GTFS-Translations has the following fields.
NEW_TRANSLATIONS_FIELDS = [
    "table_name",
    "field_name",
    "language",
    "translation",
    "record_id",
    "record_sub_id",
    "field_value",
]

# Fields whose names end with the following suffixes are translated according
# to Google translations extension.
TRANSLATABLE_FIELD_NAME_SUFFIXES = [
    "_name",
    "_desc",
    "_headsign",
    "_url",
    "_text",
    "_abbreviation",
    # Handle pathway fields "signposted_as", "reversed_signposted_as"
    # and "instructions".
    "signposted_as",
    "instructions",
]


class RecordIdHelper(object):
    """Helper object to find record_id and record_sub_id based on GTFS table
    name and its fields.
    """

    def __init__(self, table_name, field_names):
        id_and_sub_id = RECORD_ID_MAP.get(table_name)
        if id_and_sub_id is None:
            # Use the first field name that ends with _id as record_id.
            first_id = self._find_first_id(field_names)
            if first_id:
                id_and_sub_id = (first_id, None)
            else:
                id_and_sub_id = (None, None)
        self.id_and_sub_id = id_and_sub_id

    def get_record_id(self, row):
        if self.id_and_sub_id[0]:
            return row.get(self.id_and_sub_id[0])
        return None

    def get_record_sub_id(self, row):
        if self.id_and_sub_id[1]:
            return row.get(self.id_and_sub_id[1])
        return None

    def describe_ids(self):
        return 'record_id = "%s", record_sub_id = "%s"' % (
            self.id_and_sub_id[0] or "",
            self.id_and_sub_id[1] or "",
        )

    @staticmethod
    def _find_first_id(field_names):
        for field_name in field_names:
            if field_name.endswith("_id"):
                return field_name
        return None


def read_first_available_value(filename, field_name):
    """Reads the first assigned value of the given field in the CSV table.
    """
    if not os.path.exists(filename):
        return None
    with open(filename, "rt", encoding="utf-8") as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            value = row.get(field_name)
            if value:
                return value
    return None


def is_translatable_field(field):
    for suffix in TRANSLATABLE_FIELD_NAME_SUFFIXES:
        if field.endswith(suffix):
            return True
    return False


def any_translatable_field(fields):
    return any(is_translatable_field(field) for field in fields)


class OldTranslations(object):
    """Reads all old translations and keeps them for further usage.
    """

    def __init__(self, src_dir):
        self.src_dir = src_dir
        self._find_feed_language()
        self._read_translations()
        self._find_context_dependent_names()

    def _find_feed_language(self):
        """Find feed language based specified feed_info.txt or agency.txt.
        """
        self.feed_language = read_first_available_value(
            os.path.join(self.src_dir, "feed_info.txt"), "feed_lang"
        ) or read_first_available_value(
            os.path.join(self.src_dir, "agency.txt"), "agency_lang"
        )
        if not self.feed_language:
            raise Exception(
                "Cannot find feed language in feed_info.txt and agency.txt"
            )
        print("\tfeed language: %s" % self.feed_language)

    def _read_translations(self):
        """Read from the old translations.txt.
        """
        print("Reading original translations")
        self.translations_map = {}
        n_translations = 0
        with open(
            os.path.join(self.src_dir, "translations.txt"),
            "rt",
            encoding="utf-8",
        ) as csvfile:
            reader = csv.DictReader(csvfile)
            for row in reader:
                self.translations_map.setdefault(row["trans_id"], {})[
                    row["lang"]
                ] = row["translation"]
                n_translations += 1
        print("\ttotal original translations: %s" % n_translations)

    def _find_context_dependent_names(self):
        """Finds texts whose translation depends on context.

        Example.
          Here the word "Palacio" is translated from Spanish to English in
          multiple ways. Feed language is es (Spanish).

          trans_id,lang,translation
          stop-name-1,es,Palacio
          stop-name-1,en,Palace
          headsign-1,es,Palacio
          headsign-1,en,To Palace
        """
        n_occurences_of_original = {}
        for trans_id, translations in list(self.translations_map.items()):
            try:
                original_name = translations[self.feed_language]
            except KeyError:
                raise Exception(
                    "No translation in feed language for %s, available: %s"
                    % (trans_id, translations)
                )
            n_occurences_of_original[original_name] = (
                n_occurences_of_original.get(original_name, 0) + 1
            )

        self.context_dependent_names = set(
            name
            for name, occur in list(n_occurences_of_original.items())
            if occur > 1
        )
        print(
            "Total context-dependent translations: %d"
            % len(self.context_dependent_names)
        )


class TranslationsConverter(object):
    """Converts translations from the old to the new format.
    """

    def __init__(self, src_dir):
        self.src_dir = src_dir
        self.old_translations = OldTranslations(src_dir)

    def convert_translations(self, dest_dir):
        """
        Converts translations to the new format and stores at dest_dir.
        """
        if not os.path.isdir(dest_dir):
            os.makedirs(dest_dir)
        total_translation_rows = 0
        with open(os.path.join(dest_dir, "translations.txt"), "w") as out_file:
            writer = csv.DictWriter(
                out_file, fieldnames=NEW_TRANSLATIONS_FIELDS
            )
            writer.writeheader()
            for filename in sorted(os.listdir(self.src_dir)):
                if not (
                    filename.endswith(".txt")
                    and os.path.isfile(os.path.join(self.src_dir, filename))
                ):
                    print("Skipping %s" % filename)
                    continue
                table_name = filename[: -len(".txt")]
                if table_name == "translations":
                    continue
                total_translation_rows += self._translate_table(
                    dest_dir, table_name, writer
                )
        print("Total translation rows: %s" % total_translation_rows)

    def _translate_table(self, dest_dir, table_name, translations_writer):
        """
        Converts translations to the new format for a single table.
        """
        in_filename = os.path.join(self.src_dir, "%s.txt" % table_name)
        if not os.path.exists(in_filename):
            raise Exception("No %s" % table_name)

        out_filename = os.path.join(dest_dir, "%s.txt" % table_name)
        with open(in_filename, "rt", encoding="utf-8") as in_file:
            reader = csv.DictReader(in_file)
            if not reader.fieldnames or not any_translatable_field(
                reader.fieldnames
            ):
                print("Copying %s with no translatable columns" % table_name)
                shutil.copy(in_filename, out_filename)
                return 0
            table_translator = TableTranslator(
                table_name,
                reader.fieldnames,
                self.old_translations,
                translations_writer,
            )
            with open(out_filename, "w") as out_file:
                writer = csv.DictWriter(out_file, fieldnames=reader.fieldnames)
                writer.writeheader()
                for row in reader:
                    writer.writerow(table_translator.translate_row(row))

            table_translator.write_for_field_values()
            print(
                "\ttranslation rows: %s"
                % table_translator.total_translation_rows
            )
            return table_translator.total_translation_rows


class TableTranslator(object):
    """Translates a given GTFS table.
    """

    def __init__(
        self, table_name, field_names, old_translations, translations_writer
    ):
        self.table_name = table_name
        self.old_translations = old_translations
        self.translations_writer = translations_writer
        self.record_id_helper = RecordIdHelper(table_name, field_names)
        self.total_translation_rows = 0
        # stop_times.txt and trips.txt usually have a lot of repeated
        # headsigns, so it is better to use field_value than record_id
        # and record_sub_id.  However, we will fallback to
        # record_id+sub_id if the translation is context-dependent,
        # e.g., the same trip_headsign is translated differently for
        # different trips.
        self.table_uses_record_id = table_name not in ("stop_times", "trips")
        self.translations_for_values = {}

        print(
            "Translating %s by %s"
            % (
                table_name,
                self.record_id_helper.describe_ids()
                if self.table_uses_record_id
                else "field_name",
            )
        )

    def translate_row(self, row):
        table_name = self.table_name
        feed_language = self.old_translations.feed_language
        translations_map = self.old_translations.translations_map
        context_dependent_names = self.old_translations.context_dependent_names
        out_row = row
        for field_name, field_value in list(row.items()):
            if not is_translatable_field(field_name):
                continue
            field_translations = translations_map.get(field_value)
            if not field_translations:
                continue
            value_in_feed_lang = field_translations[feed_language]
            out_row[field_name] = value_in_feed_lang
            # If translation depends on the context, then always use record_id.
            use_record_id = (
                self.table_uses_record_id
                or value_in_feed_lang in context_dependent_names
            )
            record_id = self.record_id_helper.get_record_id(row)
            record_sub_id = self.record_id_helper.get_record_sub_id(row)
            for language, translation in list(field_translations.items()):
                if language == feed_language:
                    continue
                if use_record_id:
                    self._write_translation_row(
                        {
                            "table_name": table_name,
                            "field_name": field_name,
                            "language": language,
                            "translation": translation,
                            "record_id": record_id,
                            "record_sub_id": record_sub_id,
                        }
                    )
                else:
                    self.translations_for_values[
                        (field_name, language, value_in_feed_lang)
                    ] = translation
        return out_row

    def write_for_field_values(self):
        for ((field_name, language, field_value), translation) in list(
            self.translations_for_values.items()
        ):
            self._write_translation_row(
                {
                    "table_name": self.table_name,
                    "field_name": field_name,
                    "language": language,
                    "translation": translation,
                    "field_value": field_value,
                }
            )

    def _write_translation_row(self, row):
        self.translations_writer.writerow(row)
        self.total_translation_rows += 1


def main():
    if len(sys.argv) < 2:
        print(
            "usage: upgrade_translations.py [SRC GTFS DIR] [DEST GTFS DIR]",
            file=sys.stderr,
        )
        sys.exit(1)

    src_dir = os.path.normpath(sys.argv[1])
    if len(sys.argv) >= 3:
        dest_dir = sys.argv[2]
    else:
        dest_dir = "%s_upgraded" % src_dir

    print("Upgrading translations")
    print("\tsource directory: %s" % src_dir)
    print("\tdestination directory: %s" % dest_dir)

    TranslationsConverter(src_dir).convert_translations(dest_dir)
    print("Done!")


if __name__ == "__main__":
    main()