mirror of
https://github.com/onionshare/onionshare.git
synced 2024-10-01 01:35:40 -04:00
107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
# Copyright 2012-2013, Damian Johnson and The Tor Project
|
|
# See LICENSE for licensing information
|
|
|
|
"""
|
|
Toolkit for exporting descriptors to other formats.
|
|
|
|
**Module Overview:**
|
|
|
|
::
|
|
|
|
export_csv - Exports descriptors to a CSV
|
|
export_csv_file - Writes exported CSV output to a file
|
|
"""
|
|
|
|
import cStringIO
|
|
import csv
|
|
|
|
import stem.descriptor
|
|
import stem.prereq
|
|
|
|
|
|
class _ExportDialect(csv.excel):
|
|
lineterminator = '\n'
|
|
|
|
|
|
def export_csv(descriptors, included_fields = (), excluded_fields = (), header = True):
|
|
"""
|
|
Provides a newline separated CSV for one or more descriptors. If simply
|
|
provided with descriptors then the CSV contains all of its attributes,
|
|
labeled with a header row. Either 'included_fields' or 'excluded_fields' can
|
|
be used for more granular control over its attributes and the order.
|
|
|
|
:param Descriptor,list descriptors: either a
|
|
:class:`~stem.descriptor.Descriptor` or list of descriptors to be exported
|
|
:param list included_fields: attributes to include in the csv
|
|
:param list excluded_fields: attributes to exclude from the csv
|
|
:param bool header: if **True** then the first line will be a comma separated
|
|
list of the attribute names (**only supported in python 2.7 and higher**)
|
|
|
|
:returns: **str** of the CSV for the descriptors, one per line
|
|
:raises: **ValueError** if descriptors contain more than one descriptor type
|
|
"""
|
|
|
|
output_buffer = cStringIO.StringIO()
|
|
export_csv_file(output_buffer, descriptors, included_fields, excluded_fields, header)
|
|
return output_buffer.getvalue()
|
|
|
|
|
|
def export_csv_file(output_file, descriptors, included_fields = (), excluded_fields = (), header = True):
|
|
"""
|
|
Similar to :func:`stem.descriptor.export.export_csv`, except that the CSV is
|
|
written directly to a file.
|
|
|
|
:param file output_file: file to be written to
|
|
:param Descriptor,list descriptors: either a
|
|
:class:`~stem.descriptor.Descriptor` or list of descriptors to be exported
|
|
:param list included_fields: attributes to include in the csv
|
|
:param list excluded_fields: attributes to exclude from the csv
|
|
:param bool header: if **True** then the first line will be a comma separated
|
|
list of the attribute names (**only supported in python 2.7 and higher**)
|
|
|
|
:returns: **str** of the CSV for the descriptors, one per line
|
|
:raises: **ValueError** if descriptors contain more than one descriptor type
|
|
"""
|
|
|
|
if isinstance(descriptors, stem.descriptor.Descriptor):
|
|
descriptors = (descriptors,)
|
|
|
|
if not descriptors:
|
|
return
|
|
|
|
descriptor_type = type(descriptors[0])
|
|
descriptor_type_label = descriptor_type.__name__
|
|
included_fields = list(included_fields)
|
|
|
|
# If the user didn't specify the fields to include then export everything,
|
|
# ordered alphabetically. If they did specify fields then make sure that
|
|
# they exist.
|
|
|
|
desc_attr = sorted(vars(descriptors[0]).keys())
|
|
|
|
if included_fields:
|
|
for field in included_fields:
|
|
if not field in desc_attr:
|
|
raise ValueError("%s does not have a '%s' attribute, valid fields are: %s" % (descriptor_type_label, field, ", ".join(desc_attr)))
|
|
else:
|
|
included_fields = [attr for attr in desc_attr if not attr.startswith('_')]
|
|
|
|
for field in excluded_fields:
|
|
try:
|
|
included_fields.remove(field)
|
|
except ValueError:
|
|
pass
|
|
|
|
writer = csv.DictWriter(output_file, included_fields, dialect = _ExportDialect(), extrasaction='ignore')
|
|
|
|
if header and stem.prereq.is_python_27():
|
|
writer.writeheader()
|
|
|
|
for desc in descriptors:
|
|
if not isinstance(desc, stem.descriptor.Descriptor):
|
|
raise ValueError("Unable to export a descriptor CSV since %s is not a descriptor." % type(desc).__name__)
|
|
elif descriptor_type != type(desc):
|
|
raise ValueError("To export a descriptor CSV all of the descriptors must be of the same type. First descriptor was a %s but we later got a %s." % (descriptor_type_label, type(desc)))
|
|
|
|
writer.writerow(vars(desc))
|