import os
from contextlib import contextmanager
from unittest import TestCase
from dynd import nd

from blaze.data import Concat, CSV, Stack
from blaze.utils import filetexts




class Test_Files(TestCase):
    data = {'a.csv': '1,1\n2,2',
            'b.csv': '3,3\n4,4\n5,5',
            'c.csv': '6,6\n7,7'}
    def test_Concat(self):
        with filetexts(self.data) as filenames:
            descriptors = [CSV(fn, schema='2 * int32')
                            for fn in sorted(filenames)]
            dd = Concat(descriptors)

            self.assertEqual(str(dd.schema), '2 * int32')
            self.assertEqual(str(dd.dshape), 'var * 2 * int32')

            expected = ((1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6), (7, 7))

            self.assertEqual(tuple(dd), expected)

            result = dd.as_dynd()
            expected2 = nd.array(expected, dtype='int32')
            self.assertEqual(nd.as_py(result),
                             nd.as_py(expected2))

            self.assertEqual(tuple(dd), expected)
            self.assertEqual(tuple(dd), expected)  # Not one use only

            chunks = list(dd.chunks())
            assert all(isinstance(chunk, nd.array) for chunk in chunks)

            self.assertEqual(dd[[0, 2], 0], (1, 3))
            self.assertEqual(dd[2, [1, 0]], (3, 3))


class Test_Stack(TestCase):
    data = {'a.csv': '1,1\n2,2',
            'b.csv': '3,3\n4,4',
            'c.csv': '5,5\n6,6'}
    def test_Stack(self):
        with filetexts(self.data) as filenames:
            descriptors = [CSV(fn, schema='2 * int32')
                            for fn in sorted(filenames)]
            dd = Stack(descriptors)
            self.assertEqual(dd.dshape, 3 * descriptors[0].dshape)

            expected = (((1, 1), (2, 2)),
                        ((3, 3), (4, 4)),
                        ((5, 5), (6, 6)))

            self.assertEqual(dd.as_py(), expected)

            result = dd.as_dynd()
            expected2 = nd.array(expected, dtype='int32')
            self.assertEqual(nd.as_py(result),
                             nd.as_py(expected2))

            self.assertEqual(tuple(dd), expected)
            self.assertEqual(tuple(dd), expected)  # Not one use only

            chunks = dd.chunks()
            assert all(isinstance(chunk, nd.array) for chunk in chunks)

            self.assertEqual(dd[[0, 2], 0, 0], (1, 5))
            self.assertEqual(dd[0], ((1, 1), (2, 2)))
            self.assertEqual(dd[0, :, [1]], ((1,), (2,)))

@@ -1,12 +1,78 @@
+from __future__ import absolute_import, division, print_function%0A%0A
 import os%0Afr
@@ -156,16 +156,28 @@
 mport nd
+%0Aimport json
 %0A%0Afrom b
@@ -211,16 +211,32 @@
 V, Stack
+, JSON_Streaming
 %0Afrom bl
@@ -2583,12 +2583,897 @@
 1,), (2,)))%0A
+%0A%0A%0Aclass Test_Stack_JSON(TestCase):%0A    data = %7B'a.csv': %5B%7B'x':  1, 'y':  2%7D, %7B'x':  3, 'y':  4%7D%5D,%0A            'b.csv': %5B%7B'x':  5, 'y':  6%7D, %7B'x':  7, 'y':  8%7D%5D,%0A            'c.csv': %5B%7B'x':  9, 'y': 10%7D, %7B'x': 11, 'y': 12%7D%5D%7D%0A%0A    text = dict((fn, '%5Cn'.join(map(json.dumps, dicts)))%0A                    for fn, dicts in data.items())%0A    def test_Stack(self):%0A        with filetexts(self.text) as filenames:%0A            descriptors = %5BJSON_Streaming(fn, schema='%7Bx: int32, y: int32%7D')%0A                            for fn in sorted(filenames)%5D%0A            dd = Stack(descriptors)%0A%0A            expected = (((1,  2), ( 3,  4)),%0A                        ((5,  6), ( 7,  8)),%0A                        ((9, 10), (11, 12)))%0A%0A            self.assertEqual(dd.as_py(), expected)%0A%0A            self.assertEqual(dd%5B::2, 1, :%5D, ((3, 4), (11, 12)))%0A            self.assertEqual(dd%5B::2, 1, 'x'%5D, (3, 11))%0A

add json + stack test