From b617494754bd40f9f1c44a7598e9417d12acd54b Mon Sep 17 00:00:00 2001 From: Michael Kryukov Date: Mon, 8 Jan 2024 19:54:40 +0300 Subject: [PATCH] feat: copy sheetViews's pane from the template if present --- CHANGELOG.rst | 2 ++ tests/test_render.py | 19 +++++++++++++++++-- tests/utils.py | 8 +++++--- xlsx_streaming/render.py | 32 ++++++++++++++++++++++++++++++-- 4 files changed, 54 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 822998d..a3ee367 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -3,6 +3,8 @@ - Add support for Python 3.10 & 3.11 & 3.12 - Drop support for Python 3.6 & 3.7 +- If provided template contains ``sheetViews`` with a ``pane`` element, this + element will now be copied to the resulting Excel document. 1.1.0 (2021-06-23) diff --git a/tests/test_render.py b/tests/test_render.py index 3f141b4..41a9709 100644 --- a/tests/test_render.py +++ b/tests/test_render.py @@ -16,9 +16,10 @@ def test_rm_namespace(self): self.assertEqual(next(element_iter).tag, 'sheetPr') self.assertEqual(next(element_iter).tag, 'outlinePr') - def test_get_header_and_row_template(self): - header, template = render.get_elements_from_template(gen_xlsx_sheet()) # pylint: disable=unbalanced-tuple-unpacking + def test_get_elements_from_template(self): + header, views, template = render.get_elements_from_template(gen_xlsx_sheet()) # pylint: disable=unbalanced-tuple-unpacking self.assertTrue(header is None) + self.assertTrue(views is None) self.assertEqual(template.tag, 'row') self.assertEqual(template.get('r'), '1') element_iter = template.iter() @@ -27,6 +28,20 @@ def test_get_header_and_row_template(self): self.assertEqual(child.tag, 'c') self.assertEqual(child.get('r'), 'A1') + def test_views_from_template(self): + header, views, _ = render.get_elements_from_template( + gen_xlsx_sheet(with_header=True, with_views=True) + ) + self.assertTrue(header is not None) + self.assertTrue(views is not None) + self.assertEqual(views.tag, 'sheetViews') + self.assertEqual(len(views), 1) + self.assertEqual(len(views[0]), 1) + self.assertEqual(views[0][0].tag, 'pane') + self.assertEqual(views[0][0].get('xSplit'), None) + self.assertEqual(views[0][0].get('ySplit'), '1') + self.assertEqual(views[0][0].get('state'), 'frozen') + def test_get_column(self): self.assertEqual(render.get_column(ETree.Element('c', r='A1')), 'A') self.assertEqual(render.get_column(ETree.Element('c', r='ABC123')), 'ABC') diff --git a/tests/utils.py b/tests/utils.py index 8cc32e1..3d4bdb3 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -8,7 +8,7 @@ from xlsx_streaming import streaming -def gen_xlsx_template(with_header=False): +def gen_xlsx_template(with_header=False, with_views=False): wb = openpyxl.Workbook() rows = [[42, 'éOui>€', datetime.datetime(2012, 1, 2, 10, 10)]] if with_header: @@ -16,14 +16,16 @@ def gen_xlsx_template(with_header=False): for i, row in enumerate(rows): for j, value in enumerate(row): wb.active.cell(row=i + 1, column=j + 1).value = value + if with_views: + wb.active.freeze_panes = "A2" with tempfile.NamedTemporaryFile() as fp: openpyxl.writer.excel.save_workbook(wb, fp) fp.seek(0) return io.BytesIO(fp.read()) -def gen_xlsx_sheet(with_header=False): - xlsx_template = gen_xlsx_template(with_header=with_header) +def gen_xlsx_sheet(with_header=False, with_views=False): + xlsx_template = gen_xlsx_template(with_header=with_header, with_views=with_views) with zipfile.ZipFile(xlsx_template, mode='r') as zip_template: sheet_name = streaming.get_first_sheet_name(zip_template) return zip_template.read(sheet_name).decode('utf-8') diff --git a/xlsx_streaming/render.py b/xlsx_streaming/render.py index 5a06a31..7cc11df 100644 --- a/xlsx_streaming/render.py +++ b/xlsx_streaming/render.py @@ -32,10 +32,13 @@ def render_worksheet(rows_batches, openxml_sheet_string, encoding='utf-8'): rows_batches (iterable): each element is a list of lists containing the row values openxml_sheet_string (str): a template for the final sheet containing the header and an example row """ - header_tree, row_template = get_elements_from_template(openxml_sheet_string) + header_tree, views, row_template = get_elements_from_template(openxml_sheet_string) yield f'\n'.encode(encoding) + if views is not None: + yield ETree.tostring(views, encoding=encoding) + yield '\n'.encode(encoding) current_line = 1 @@ -289,13 +292,38 @@ def _get_header_and_row_template(tree): return None, header_tree +def _get_sheet_views(tree): + """ + Extract sheet views (potentially None) from the provided tree + args: + tree (ElementTree.Element): root element of the template + return (ElementTree.Element): + Constructed sheetViews ElementTree.Element object + """ + pane = tree.find('sheetViews/sheetView/pane') + + # Currently we only support panes with fronzen state + if pane is None or pane.get('state') != 'frozen': + return None + + sheet_views = ETree.Element('sheetViews') + sheet_view = ETree.Element('sheetView', {'workbookViewId': '0'}) + sheet_views.append(sheet_view) + pane = ETree.Element('pane', pane.attrib) + sheet_view.append(pane) + + return sheet_views + + def get_elements_from_template(openxml_sheet): tree = ETree.fromstring(openxml_sheet) rm_namespace(tree) header, row_template = _get_header_and_row_template(tree) - return header, row_template + views = _get_sheet_views(tree) + + return header, views, row_template def get_default_template(row_values, reset_memory=False):