#!/usr/bin/env python # # Copyright 2007 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. # """Bulkloader CSV reading and writing. Handle the CSV format specified in a bulkloader.yaml file. """ import codecs import cStringIO import csv import encodings import encodings.ascii import encodings.cp1252 import encodings.latin_1 import encodings.utf_8 from google.appengine.ext.bulkload import bulkloader_errors from google.appengine.ext.bulkload import connector_interface def utf8_recoder(stream, encoding): """Generator that reads an encoded stream and reencodes to UTF-8.""" for line in codecs.getreader(encoding)(stream): yield line.encode('utf-8') class UnicodeDictWriter(object): """Based on UnicodeWriter in http://docs.python.org/library/csv.html.""" def __init__(self, stream, fieldnames, encoding='utf-8', **kwds): """Initialzer. Args: stream: Stream to write to. fieldnames: Fieldnames to pass to the DictWriter. encoding: Desired encoding. kwds: Additional arguments to pass to the DictWriter. """ writer = codecs.getwriter(encoding) if (writer is encodings.utf_8.StreamWriter or writer is encodings.ascii.StreamWriter or writer is encodings.latin_1.StreamWriter or writer is encodings.cp1252.StreamWriter): self.no_recoding = True self.encoder = codecs.getencoder(encoding) self.writer = csv.DictWriter(stream, fieldnames, **kwds) else: self.no_recoding = False self.encoder = codecs.getencoder('utf-8') self.queue = cStringIO.StringIO() self.writer = csv.DictWriter(self.queue, fieldnames, **kwds) self.stream = writer(stream) def writerow(self, row): """Wrap writerow method.""" row_encoded = dict([(k, self.encoder(v)[0]) for (k, v) in row.iteritems()]) self.writer.writerow(row_encoded) if self.no_recoding: return data = self.queue.getvalue() data = data.decode('utf-8') self.stream.write(data) self.queue.truncate(0) class CsvConnector(connector_interface.ConnectorInterface): """Read/write a (possibly encoded) CSV file.""" @classmethod def create_from_options(cls, options, name): """Factory using an options dictionary. Args: options: Dictionary of options: columns: 'from_header' or blank. column_list: overrides columns specifically. encoding: encoding of the file. e.g. 'utf-8' (default), 'windows-1252'. skip_import_header_row: True to ignore the header line on import. Defaults False, except must be True if columns=from_header. print_export_header_row: True to print a header line on export. Defaults to False except if columns=from_header. import_options: Other kwargs to pass in, like "dialect". export_options: Other kwargs to pass in, like "dialect". name: The name of this transformer, for use in error messages. Returns: CsvConnector object described by the specified options. Raises: InvalidConfiguration: If the config is invalid. """ column_list = options.get('column_list', None) columns = None if not column_list: columns = options.get('columns', 'from_header') if columns != 'from_header': raise bulkloader_errors.InvalidConfiguration( 'CSV columns must be "from_header", or a column_list ' 'must be specified. (In transformer name %s.)' % name) csv_encoding = options.get('encoding', 'utf-8') skip_import_header_row = options.get('skip_import_header_row', columns == 'from_header') if columns == 'from_header' and not skip_import_header_row: raise bulkloader_errors.InvalidConfiguration( 'When CSV columns are "from_header", the header row must always ' 'be skipped. (In transformer name %s.)' % name) print_export_header_row = options.get('print_export_header_row', columns == 'from_header') import_options = options.get('import_options', {}) export_options = options.get('export_options', {}) return cls(columns, column_list, skip_import_header_row, print_export_header_row, csv_encoding, import_options, export_options) def __init__(self, columns, column_list, skip_import_header_row, print_export_header_row, csv_encoding=None, import_options=None, export_options=None): """Initializer. Args: columns: 'from_header' or blank column_list: overrides columns specifically. skip_import_header_row: True to ignore the header line on import. Defaults False, except must be True if columns=from_header. print_export_header_row: True to print a header line on export. Defaults to False except if columns=from_header. csv_encoding: encoding of the file. import_options: Other kwargs to pass in, like "dialect". export_options: Other kwargs to pass in, like "dialect". """ self.columns = columns self.from_header = (columns == 'from_header') self.column_list = column_list self.skip_import_header_row = skip_import_header_row self.print_export_header_row = print_export_header_row self.csv_encoding = csv_encoding self.dict_generator = None self.output_stream = None self.csv_writer = None self.bulkload_state = None self.import_options = import_options or {} self.export_options = export_options or {} def generate_import_record(self, filename, bulkload_state): """Generator, yields dicts for nodes found as described in the options. Args: filename: Filename to read. bulkload_state: Passed bulkload_state. Yields: Neutral dict, one per row in the CSV file. """ self.bulkload_state = bulkload_state input_stream = open(filename) input_stream = utf8_recoder(input_stream, self.csv_encoding) self.dict_generator = csv.DictReader(input_stream, self.column_list, **self.import_options) discard_line = self.skip_import_header_row and not self.from_header line_number = 0 for input_dict in self.dict_generator: line_number = line_number + 1 if discard_line: discard_line = False continue decoded_dict = {} for key, value in input_dict.iteritems(): if key == None: raise bulkloader_errors.InvalidImportData( 'Got more values in row than headers on line %d.' % (line_number)) if not self.column_list: key = unicode(key, 'utf-8') if value: value = unicode(value, 'utf-8') decoded_dict[key] = value yield decoded_dict def initialize_export(self, filename, bulkload_state): """Initialize the output file. Args: filename: Filename to write. bulkload_state: Passed bulkload_state. """ self.bulkload_state = bulkload_state self.output_stream = open(filename, 'wb') def __initialize_csv_writer(self, dictionary): """Actual initialization, happens on the first entity being written.""" write_header = self.print_export_header_row if self.from_header: export_column_list = tuple(dictionary) else: export_column_list = self.column_list self.csv_writer = UnicodeDictWriter(self.output_stream, export_column_list, self.csv_encoding, **self.export_options) if write_header: self.csv_writer.writerow(dict(zip(export_column_list, export_column_list))) def write_dict(self, dictionary): """Write one record for the specified entity.""" if not self.csv_writer: self.__initialize_csv_writer(dictionary) self.csv_writer.writerow(dictionary) def finalize_export(self): self.output_stream.close()