diff --git a/rows/__init__.py b/rows/__init__.py index 897df4e8..a5752f43 100644 --- a/rows/__init__.py +++ b/rows/__init__.py @@ -58,5 +58,9 @@ except ImportError: pass +try: + from rows.plugins._pandas import import_from_pandas, export_to_pandas +except ImportError: + pass __version__ = '0.2.0-dev' diff --git a/rows/plugins/_pandas.py b/rows/plugins/_pandas.py new file mode 100644 index 00000000..fe3aff05 --- /dev/null +++ b/rows/plugins/_pandas.py @@ -0,0 +1,60 @@ +# coding: utf-8 + +# Copyright 2014-2015 Álvaro Justen +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +import datetime + +import pandas + +from rows.plugins.utils import create_table + +DATE_FORMAT = "%Y-%m-%d %H:%M:%S" + +def import_from_pandas(data_frame, *args, **kwargs): + meta = {'imported_from': 'pandas', 'filename': 'DataFrame', } + return create_table(_dataframe_generator(data_frame), meta=meta, *args, + **kwargs) + +def _dataframe_generator(data_frame): + yield list(data_frame) + + for _, row in data_frame.iterrows(): + yield dataframe_row_to_list(row) + +def dataframe_row_to_list(row): + result = [] + for element_index, element in enumerate(row): + #Problem importing pandas.tslib.Timestamp or pandas.Timestamp + if isinstance(element, pandas.tslib.Timestamp): + date_string = element.strftime(DATE_FORMAT) + if date_string.endswith("00:00:00"): + value = datetime.datetime.strptime(date_string, + DATE_FORMAT).date() + else: + value = element.to_datetime() + result.append(value) + return result + + +def export_to_pandas(table_obj): + data_frame = pandas.DataFrame(_generator_table(table_obj), + columns=table_obj.field_names) + + return data_frame + +def _generator_table(table_obj): + for row in table_obj: + yield list(row) diff --git a/setup.py b/setup.py index b5a57fa8..f8d304f6 100644 --- a/setup.py +++ b/setup.py @@ -24,7 +24,8 @@ 'html': ['lxml'], # apt: libxslt-dev libxml2-dev 'ods': ['lxml'], 'xls': ['xlrd', 'xlwt'], - 'xlsx': ['openpyxl'], } + 'xlsx': ['openpyxl'], + 'pandas': ['pandas'], } EXTRA_REQUIREMENTS['all'] = sum(EXTRA_REQUIREMENTS.values(), []) INSTALL_REQUIREMENTS = EXTRA_REQUIREMENTS['csv'] LONG_DESCRIPTION = ''' diff --git a/tests/tests_plugin_pandas.py b/tests/tests_plugin_pandas.py new file mode 100644 index 00000000..9d140b73 --- /dev/null +++ b/tests/tests_plugin_pandas.py @@ -0,0 +1,56 @@ +# coding: utf-8 + +# Copyright 2014-2015 Álvaro Justen +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +import unittest +import mock + +import pandas + +import rows.plugins._pandas +import rows.plugins.csv +import utils + +class PluginPandasTestCase(utils.RowsTestMixIn, unittest.TestCase): + + plugin_name = 'pandas' + data_frame = pandas.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]], + columns=['A', 'B', 'C', 'D']) + + def test_imports(self): + self.assertIs(rows.import_from_pandas, + rows.plugins._pandas.import_from_pandas) + + @mock.patch('rows.plugins._pandas.create_table') + def test_import_from_pandas_uses_create_table(self, mocked_create_table): + mocked_create_table.return_value = 101 + kwargs = {'encoding': 'test', 'some_key': 123, 'other': 456, } + result = rows.import_from_pandas(self.data_frame, **kwargs) + self.assertTrue(mocked_create_table.called) + self.assertEqual(mocked_create_table.call_count, 1) + self.assertEqual(result, 101) + + call = mocked_create_table.call_args + kwargs['meta'] = {'imported_from': 'pandas', 'filename': 'DataFrame'} + self.assertEqual(call[1], kwargs) + + @mock.patch('rows.plugins._pandas.pandas.DataFrame') + def test_export_to_pandas_uses_data_frame(self, mocked_data_frame): + mocked_data_frame.return_value = 101 + result = rows.plugins._pandas.export_to_pandas(utils.table) + self.assertTrue(mocked_data_frame.called) + self.assertTrue(mocked_data_frame.call_count, 1) + self.assertEqual(result, 101) +