Permalink
Browse files

Add dictionary support to csvpy. Fix obscure CSVKitDictReader bugs.

  • Loading branch information...
1 parent 9a62bdc commit 19b13ea4bef8cd61fa1015cea04d065764b9dcfd @onyxfish onyxfish committed Aug 15, 2012
Showing with 46 additions and 9 deletions.
  1. +2 −4 csvkit/__init__.py
  2. +3 −0 csvkit/cli.py
  3. +7 −1 csvkit/unicsv.py
  4. +11 −3 csvkit/utilities/csvpy.py
  5. +9 −1 docs/scripts/csvpy.rst
  6. +14 −0 tests/test_unicsv.py
View
@@ -10,8 +10,7 @@
class CSVKitReader(unicsv.UnicodeCSVReader):
"""
- A unicode-aware CSV reader. Currently adds nothing to :class:`csvkit.unicsv.UnicodeCSVReader`, but might
- someday.
+ A unicode-aware CSV reader. Currently adds nothing to :class:`csvkit.unicsv.UnicodeCSVReader`, but might someday.
"""
pass
@@ -48,8 +47,7 @@ def writerows(self, rows):
class CSVKitDictReader(unicsv.UnicodeCSVDictReader):
"""
- A unicode-aware CSV DictReader. Currently adds nothing to :class:`csvkit.unicsv.UnicodeCSVWriter`, but might
- someday.
+ A unicode-aware CSV DictReader. Currently adds nothing to :class:`csvkit.unicsv.UnicodeCSVWriter`, but might someday.
"""
pass
View
@@ -38,6 +38,9 @@ def __getattr__(self, name):
return getattr(self.f, name)
+ def __iter__(self):
+ return self
+
def close(self):
self.f.close()
self.f = None
View
@@ -102,8 +102,14 @@ class UnicodeCSVDictReader(csv.DictReader):
of :func:`csv.reader`.
"""
def __init__(self, f, fieldnames=None, restkey=None, restval=None, *args, **kwargs):
+ reader = UnicodeCSVReader(f, *args, **kwargs)
+
+ if 'encoding' in kwargs:
+ kwargs.pop('encoding')
+
csv.DictReader.__init__(self, f, fieldnames, restkey, restval, *args, **kwargs)
- self.reader = UnicodeCSVReader(f, *args, **kwargs)
+
+ self.reader = reader
class UnicodeCSVDictWriter(csv.DictWriter):
"""
@@ -1,6 +1,6 @@
#!/usr/bin/env python
-from csvkit import CSVKitReader
+from csvkit import CSVKitReader, CSVKitDictReader
from csvkit.cli import CSVFileType, CSVKitUtility
class CSVPy(CSVKitUtility):
@@ -10,13 +10,21 @@ class CSVPy(CSVKitUtility):
def add_arguments(self):
self.argparser.add_argument('file', metavar="FILE", type=CSVFileType(),
help='The CSV file to operate on.')
+ self.argparser.add_argument('--dict', dest='as_dict', action='store_true',
+ help='Use CSVKitDictReader instead of CSVKitReader.')
def main(self):
# Attempt reading filename, will cause lazy loader to access file and raise error if it does not exist
filename = self.args.file.name
- welcome_message = 'Welcome! "%s" has been loaded in a CSVKitReader object named "reader".' % filename
- reader = CSVKitReader(self.args.file, **self.reader_kwargs)
+ if self.args.as_dict:
+ reader_class = CSVKitDictReader
+ else:
+ reader_class = CSVKitReader
+
+ reader = reader_class(self.args.file, **self.reader_kwargs)
+
+ welcome_message = 'Welcome! "%s" has been loaded in a %s object named "reader".' % (filename, reader_class.__name__)
try:
from IPython.frontend.terminal.embed import InteractiveShellEmbed
@@ -5,7 +5,7 @@ csvpy
Description
===========
-Loads a CSV file into a :doc:`CSVKitReader </api/csvkit>` object and then drops into a Python shell so the user can inspect the data however they see fit::
+Loads a CSV file into a :class:`csvkit.CSVKitReader` object and then drops into a Python shell so the user can inspect the data however they see fit::
usage: csvpy [-h] [-d DELIMITER] [-t] [-q QUOTECHAR] [-u {0,1,2,3}] [-b]
[-p ESCAPECHAR] [-z MAXFIELDSIZE] [-e ENCODING] [-v]
@@ -18,6 +18,7 @@ Loads a CSV file into a :doc:`CSVKitReader </api/csvkit>` object and then drops
optional arguments:
-h, --help show this help message and exit
+ --dict Use CSVKitDictReader instead of CSVKitReader.
This utility will automatically use the IPython shell if it is installed, otherwise it will use the running Python shell.
@@ -37,3 +38,10 @@ Basic use::
>>> reader.next()
[u'a', u'b', u'c']
+As a dictionary::
+
+ $ csvpy --dict examples/dummy.csv -v
+ Welcome! "examples/dummy.csv" has been loaded in a CSVKitDictReader object named "reader".
+ >>> reader.next()
+ {u'a': u'1', u'c': u'3', u'b': u'2'}
+
View
@@ -114,6 +114,20 @@ def test_reader(self):
u'c': u'3'
})
+ def test_latin1(self):
+ with open('examples/test_latin1.csv') as f:
+ reader = unicsv.UnicodeCSVDictReader(f, encoding='latin1')
+ self.assertEqual(reader.next(), {
+ u'a': u'1',
+ u'b': u'2',
+ u'c': u'3'
+ })
+ self.assertEqual(reader.next(), {
+ u'a': u'4',
+ u'b': u'5',
+ u'c': u'©'
+ })
+
class TestUnicodeCSVDictWriter(unittest.TestCase):
def setUp(self):
self.output = StringIO()

0 comments on commit 19b13ea

Please sign in to comment.