[tor-commits] [stem/master] First draft of torexport, 3 functions defined to handle different use cases.

[atagar at torproject.org]

commit 508373f1a7824b1ff9b86a597a718918d4250e24
Author: Erik <eislo at wesleyan.edu>
Date:   Thu Jul 12 14:13:59 2012 -0400

    First draft of torexport, 3 functions defined to handle different use cases.
    
    First draft of torexport, 3 functions defined to handle different use cases.
    csv_file_exp() takes a list of descriptors and exports their attributes as a
    csv file. descriptors_csv_exp() takes a list of descriptors and returns a
    generator.  Each iteration, this function returns a line of csv representing
    the attributes of a single descriptor object.  descriptor_csv_exp() returns
    a single line of csv based on the attributes of a sole descriptor object.
---
 stem/descriptor/export.py |   89 +++++++++++++++++++++++++++++++++++++++++++++
 1 files changed, 89 insertions(+), 0 deletions(-)

diff --git a/stem/descriptor/export.py b/stem/descriptor/export.py
new file mode 100644
index 0000000..2859fee
--- /dev/null
+++ b/stem/descriptor/export.py
@@ -0,0 +1,89 @@
+import os, csv, sets, cStringIO
+
+
+def descriptor_csv_exp(descr, incl_fields=[], excl_fields=[]):
+  """
+  Takes a single descriptor object, puts it in a list, and passes it to
+  descriptors_csv_exp to build a csv.
+
+  :param object descr: single descriptor object to export as csv.
+  """
+  descriptor = [descr]
+  for desc in descriptors_csv_exp(descriptor, incl_fields, excl_fields):
+    return desc
+
+
+def descriptors_csv_exp(descrs, incl_fields=[], excl_fields=[], head=False):
+  """
+  Builds a csv file based on attributes of descriptors.
+
+  :param list descrs: List of descriptor objects to export as a csv line.
+  :param list incl_fields: list of attribute fields to include in the csv line.
+  :param list excl_fields: list of attribute fields to exclude from csv line.
+  :param bool head: whether or not a header is requested; shouldn't be needed
+    outside of csv_file_exp's call of this function.
+
+  :returns: generator for csv strings, one line per descr object.
+  """
+  
+  _temp_file = cStringIO.StringIO()
+
+  first = True
+
+  for desc in descrs:
+    attr = vars(desc)
+
+    # Defining incl_fields and the dwriter object requires having access
+    # to a descriptor object.
+    if first:
+      # define incl_fields, 4 cases where final case is incl_fields already
+      # defined and excl_fields left blank, so no action is necessary.
+      if not incl_fields and excl_fields:
+        _incl = sets.Set(attr.keys())
+        incl_fields = list(_incl.difference(excl_fields))
+
+      elif not incl_fields and not excl_fields:
+        incl_fields = attr.keys()
+
+      elif incl_fields and excl_fields:
+        _incl = sets.Set(incl_fields)
+        incl_fields = list(_incl.difference(excl_fields))
+
+      dwriter = csv.DictWriter(_temp_file, incl_fields)
+      first = False
+
+      if head:
+        dwriter.writeheader()
+
+    # Need to remove fields that aren't wanted for dwriter.
+    final = {}
+    for at in attr:
+      if at in incl_fields:
+        final[at] = attr[at]
+
+    dwriter.writerow(final)
+    yield _temp_file.getvalue()
+    
+    # Clear cString wrapper for new descriptor.
+    _temp_file.reset()
+    _temp_file.truncate()
+
+      
+  _temp_file.close()
+
+def csv_file_exp(descrs, doc_loc, header=True, incl_f=[], excl_f=[]):
+  """
+  Writes descriptor attributes to a csv file on disk.
+
+  :param list descrs: descriptor objects with attributes to export as csv file.
+  :param str doc_loc: location and file name for csv file to be written to.
+    This overwrites existing files.
+  :param bool header: defaults to true, determines if doc will have a header row.
+  :param list incl_f: list of attribute fields to include in the csv line.
+  :param list excl_f: list of attribute fields to exclude from csv line.
+  """
+  doc = open(doc_loc, 'w')
+
+  for line in descriptors_csv_exp(descrs, incl_fields=incl_f, excl_fields=excl_f, head=header):
+    doc.write(line)
+