1import pyexcel as pe
2from _compact import OrderedDict
3
4from nose.tools import eq_, raises
5
6
7class Attributable:
8    def __init__(self, adict):
9        self.mydict = adict
10
11    def __getattr__(self, field):
12        return self.mydict[field]
13
14
15class Objects:
16    def __init__(self):
17        self.objs = None
18
19    def bulk_create(self, objs, batch_size):
20        self.objs = objs
21        self.batch_size = batch_size
22
23    def all(self):
24        return [Attributable(o) for o in self.objs]
25
26
27class Field:
28    def __init__(self, name):
29        self.attname = name
30
31
32class Meta:
33    instance = 1
34
35    def __init__(self):
36        self.model_name = "Sheet%d" % Meta.instance
37        self.concrete_fields = []
38        Meta.instance = Meta.instance + 1
39
40    def update(self, data):
41        for f in data:
42            self.concrete_fields.append(Field(f))
43
44
45class FakeDjangoModel:
46    def __init__(self, model_name=None):
47        self.objects = Objects()
48        self._meta = Meta()
49        if model_name:
50            self._meta.model_name = model_name
51
52    def __call__(self, **keywords):
53        return keywords
54
55    def save(self):
56        pass
57
58
59class TestVerticalSheet:
60    def setUp(self):
61        self.data = [["X", 1, 4], ["Y", 2, 5], ["Z", 3, 6]]
62        self.result = [{"Y": 2, "X": 1, "Z": 3}, {"Y": 5, "X": 4, "Z": 6}]
63
64    def test_model_save_to_django_model(self):
65        model = FakeDjangoModel()
66        pe.save_as(
67            array=self.data,
68            name_columns_by_row=0,
69            dest_model=model,
70            transpose_before=True,
71        )
72        assert model.objects.objs == self.result
73
74    def test_mapping_array(self):
75        data2 = [["A", 1, 4], ["B", 2, 5], ["C", 3, 6]]
76        mapdict = ["X", "Y", "Z"]
77        model = FakeDjangoModel()
78        pe.save_as(
79            array=data2,
80            name_columns_by_row=0,
81            dest_model=model,
82            dest_mapdict=mapdict,
83            transpose_before=True,
84        )
85        assert model.objects.objs == self.result
86
87    def test_mapping_dict(self):
88        """
89
90        for vertical sheet, first transpose it and then
91        name columns by row 0
92        """
93        data2 = [["A", 1, 4], ["B", 2, 5], ["C", 3, 6]]
94        mapdict = {"C": "Z", "A": "X", "B": "Y"}
95        model = FakeDjangoModel()
96        pe.save_as(
97            array=data2,
98            dest_model=model,
99            dest_mapdict=mapdict,
100            name_columns_by_row=0,
101            transpose_before=True,
102        )
103        eq_(model.objects.objs, self.result)
104
105
106class TestSheet:
107    def setUp(self):
108        self.data = [["X", "Y", "Z"], [1, 2, 3], [4, 5, 6]]
109        self.result = [{"Y": 2, "X": 1, "Z": 3}, {"Y": 5, "X": 4, "Z": 6}]
110
111    def test_sheet_save_to_django_model(self):
112        model = FakeDjangoModel()
113        sheet = pe.Sheet(self.data, name_columns_by_row=0)
114        sheet.save_to_django_model(model)
115        assert model.objects.objs == self.result
116
117    def test_sheet_save_to_django_model_3(self):
118        model = FakeDjangoModel()
119        sheet = pe.Sheet(self.data)
120        sheet.name_columns_by_row(0)
121
122        def wrapper(row):
123            row[0] = row[0] + 1
124            return row
125
126        sheet.save_to_django_model(model, initializer=wrapper)
127        assert model.objects.objs == [
128            {"Y": 2, "X": 2, "Z": 3},
129            {"Y": 5, "X": 5, "Z": 6},
130        ]
131
132    def test_model_save_to_django_model(self):
133        model = FakeDjangoModel()
134        pe.save_as(array=self.data, name_columns_by_row=0, dest_model=model)
135        assert model.objects.objs == self.result
136
137    def test_model_save_to_django_model_2(self):
138        model = FakeDjangoModel()
139        pe.save_as(array=self.data, dest_model=model, name_columns_by_row=0)
140        assert model.objects.objs == self.result
141
142    def test_load_sheet_from_django_model(self):
143        model = FakeDjangoModel()
144        sheet = pe.Sheet(self.data, name_columns_by_row=0)
145        sheet.save_to_django_model(model)
146        assert model.objects.objs == self.result
147        model._meta.update(["X", "Y", "Z"])
148        sheet2 = pe.get_sheet(model=model, sheet_name="test")
149        sheet2.name_columns_by_row(0)
150        assert sheet2.name == "test"
151        eq_(list(sheet2.to_records()), list(sheet.to_records()))
152
153    def test_mapping_array(self):
154        data2 = [["A", "B", "C"], [1, 2, 3], [4, 5, 6]]
155        mapdict = ["X", "Y", "Z"]
156        model = FakeDjangoModel()
157        pe.save_as(
158            array=data2,
159            name_columns_by_row=0,
160            dest_model=model,
161            dest_mapdict=mapdict,
162        )
163        assert model.objects.objs == self.result
164
165    @raises(Exception)
166    def test_mapping_array_exceptional_case(self):
167        data2 = [["A", "B", "C"], [1, 2, 3], [4, 5, 6]]
168        mapdict = ["X", "Y", "Z"]
169        model = FakeDjangoModel()
170        pe.save_as(array=data2, dest_model=model, dest_mapdict=mapdict)
171        assert model.objects.objs == self.result
172
173    def test_mapping_dict(self):
174        data2 = [["A", "B", "C"], [1, 2, 3], [4, 5, 6]]
175        mapdict = {"C": "Z", "A": "X", "B": "Y"}
176        model = FakeDjangoModel()
177        pe.save_as(
178            array=data2,
179            name_columns_by_row=0,
180            dest_model=model,
181            dest_mapdict=mapdict,
182        )
183        eq_(model.objects.objs, self.result)
184
185
186class TestBook:
187    def setUp(self):
188        self.content = OrderedDict()
189        self.content.update(
190            {"Sheet1": [[u"X", u"Y", u"Z"], [1, 4, 7], [2, 5, 8], [3, 6, 9]]}
191        )
192        self.content.update(
193            {"Sheet2": [[u"A", u"B", u"C"], [1, 4, 7], [2, 5, 8], [3, 6, 9]]}
194        )
195        self.result1 = [
196            {"Y": 4, "X": 1, "Z": 7},
197            {"Y": 5, "X": 2, "Z": 8},
198            {"Y": 6, "X": 3, "Z": 9},
199        ]
200        self.result2 = [
201            {"B": 4, "A": 1, "C": 7},
202            {"B": 5, "A": 2, "C": 8},
203            {"B": 6, "A": 3, "C": 9},
204        ]
205
206    def test_book_save_to_models(self):
207        model1 = FakeDjangoModel("Sheet1")
208        model2 = FakeDjangoModel("Sheet2")
209        book = pe.Book(self.content)
210        book.save_to_django_models([model1, model2])
211        assert model1.objects.objs == self.result1
212        assert model2.objects.objs == self.result2
213
214    @raises(AttributeError)
215    def test_book_save_to_models_with_bulk_save_false(self):
216        """
217        same to previous test but with different parameters
218        """
219        model1 = FakeDjangoModel("Sheet1")
220        model2 = FakeDjangoModel("Sheet2")
221        book = pe.Book(self.content)
222        book.save_to_django_models([model1, model2], bulk_save=False)
223
224    def test_model_save_to_models(self):
225        model = FakeDjangoModel("Sheet1")
226        data = {
227            "Sheet1": [[u"X", u"Y", u"Z"], [1, 4, 7], [2, 5, 8], [3, 6, 9]]
228        }
229        pe.save_book_as(dest_models=[model, None, None], bookdict=data)
230        assert model.objects.objs == self.result1
231
232    def test_load_book_from_django_model(self):
233        # if a book has more than one sheet
234        # and it saves to only one model, now it will fail
235        # with an exception.
236        model = FakeDjangoModel("Sheet1")
237        book = pe.Book(
238            {"Sheet1": [[u"X", u"Y", u"Z"], [1, 4, 7], [2, 5, 8], [3, 6, 9]]}
239        )
240        book.save_to_django_models([model])
241        assert model.objects.objs == self.result1
242        model._meta.update(["X", "Y", "Z"])
243        book2 = pe.get_book(models=[model])
244        assert book2[0].to_array() == book[0].to_array()
245
246    @raises(Exception)
247    def test_more_sheets_than_models(self):
248        self.content.update({"IgnoreMe": [[1, 2, 3]]})
249        model = FakeDjangoModel("Sheet1")
250        pe.save_book_as(dest_models=[model], bookdict=self.content)
251