| #!/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() |
