added stem python library

lib/stem/descriptor/export.py

@@ -0,0 +1,106 @@
+# 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))