Fixed proper encoding when opening file, limiting zipping from stdin

2 files changed, 31 insertions, 6 deletions

diff --git a/bin/q b/bin/q
index d6caf21..a7213d8 100755
--- a/bin/q
+++ b/bin/q
@@ -290,6 +290,11 @@ class BadHeaderException(Exception):
         return repr(self.msg)
 
 
+class CannotUnzipStdInException(Exception):
+
+    def __init__(self):
+        pass
+
 class EmptyDataException(Exception):
 
     def __init__(self):
@@ -714,13 +719,15 @@ class TableCreator(object):
 
                 # Check if it's standard input or a file
                 if filename == '-':
-                    f = sys.stdin
+                    codec_info = codecs.lookup(self.encoding)
+                    f  = codecs.StreamReaderWriter(sys.stdin,codec_info.streamreader,codec_info.streamwriter,None)
+                    if self.gzipped:
+                        raise CannotUnzipStdInException()
                 else:
-                    f = file(filename, 'rb')
-
-                # Wrap it with gzip decompression if needed
-                if self.gzipped or filename.endswith('.gz'):
-                    f = gzip.GzipFile(fileobj=f)
+                    if self.gzipped or filename.endswith('.gz'):
+                        f = gzip.GzipFile(fileobj=file(filename,'rb'))    
+                    else:
+                        f = codecs.open(filename, 'rb',encoding=self.encoding)
 
                 self.read_file_using_csv(f, analyze_only)
                 if not self.table_created:
@@ -1054,6 +1061,9 @@ except (UnicodeDecodeError, UnicodeError), e:
 except BadHeaderException, e:
     print >>sys.stderr, "Bad header row: %s" % e.msg
     sys.exit(35)
+except CannotUnzipStdInException,e:
+    print >>sys.stderr,"Cannot decompress standard input. Pipe the input through zcat in order to decompress."
+    sys.exit(36)
 except KeyboardInterrupt:
     print >>sys.stderr, "Interrupted"
     sys.exit(0)
diff --git a/test/test-suite b/test/test-suite
index 87692ae..b167fdf 100755
--- a/test/test-suite
+++ b/test/test-suite
@@ -158,6 +158,21 @@ class BasicTests(AbstractQTestCase):
 
         self.cleanup(tmpfile)
 
+    def test_attempt_to_unzip_stdin(self):
+        tmpfile = self.create_file_with_data(
+            '\x1f\x8b\x08\x08\xf2\x18\x12S\x00\x03xxxxxx\x003\xe42\xe22\xe62\xe12\xe52\xe32\xe7\xb2\xe0\xb2\xe424\xe0\x02\x00\xeb\xbf\x8a\x13\x15\x00\x00\x00', encoding='none')
+
+        cmd = 'cat %s | ../bin/q -z "select sum(c1),avg(c1) from -"' % tmpfile.name
+
+        retcode, o, e = run_command(cmd)
+        self.assertTrue(retcode != 0)
+        self.assertTrue(len(o) == 0)
+        self.assertTrue(len(e) == 1)
+
+        self.assertEquals(e[0],'Cannot decompress standard input. Pipe the input through zcat in order to decompress.')
+
+        self.cleanup(tmpfile)
+
     def test_delimition_mistake_with_header(self):
         tmpfile = self.create_file_with_data(sample_data_no_header)