Last active
December 10, 2015 21:48
-
-
Save dmvianna/4497752 to your computer and use it in GitHub Desktop.
Extending pandas DataFrame and Series objects [http://pandas.pydata.org/] so it properly saves multiple levels of column/row indexes and xlwt styles to .xls files.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
#ws argument is a xlwt worksheet object, as in: | |
from xlwt import Workbook | |
wb = Workbook() | |
ws = wb.add.sheet('sheet 1') | |
# Example of how to wrap xlwt.Style object | |
conv = pd.io.parsers.CellStyleConverter() | |
hstyle_dict = {"font": {"bold": True}, | |
"border": {"top": "thin", | |
"right": "thin", | |
"bottom": "thin", | |
"left": "thin"}, | |
"align": {"horiz": "center"}} | |
hstyle = conv.to_xls(hstyle_dict) | |
""" | |
import pandas as pd | |
import xlwt as xlwt | |
import xlwt.Style as Style | |
class XLtable(pd.DataFrame): | |
def __init__(self, df=None): | |
""" | |
pandas DataFrame with extra methods for setting location and xlwt style | |
in an Excel spreadsheet. | |
Parameters | |
---------- | |
df : A pandas DataFrame object | |
Example | |
------- | |
>>> xldata = XLtable(df=data) | |
""" | |
super(XLtable, self).__init__() | |
self._data = df._data | |
if isinstance(self.columns, pd.MultiIndex): | |
self.col_depth = len(self.columns.levels) | |
else: | |
self.col_depth = 1 | |
if isinstance(self.index, pd.MultiIndex): | |
self.row_depth = len(self.index.levels) | |
else: | |
self.row_depth = 1 | |
def place_index(self, ws, row=0, col=0, axis=0, style=Style.default_style): | |
""" | |
Write XLtable row or column names (indexes) into an Excel sheet | |
Parameters | |
---------- | |
ws : An xlwt Worksheet object | |
row: The starting row in Excel where the index will be placed | |
(integer, base 0, default 0) | |
col: The starting column in Excel where the index will be placed | |
(integer, base 0, default 0) | |
axis: Whether the row or column index is desired | |
(0 for row or 1 for column, default 0) | |
style: An xlwt style object | |
Example | |
------- | |
>>> XLtable.place_index(ws=Sheet1, row=0, col=0, axis=0, style=hstyle) | |
""" | |
if axis == 0: | |
depth = self.row_depth | |
index = self.index | |
elif axis == 1: | |
depth = self.col_depth | |
index = self.columns | |
else: | |
raise ValueError("XLTable has only two axis (0 or 1)") | |
if depth == 1: | |
if axis == 0: | |
for i in range(len(index)): | |
ws.row(row + 1 + i).write(col, index[i], style) | |
elif axis == 1: | |
for i in range(len(index)): | |
ws.row(row).write(col + 1 + i, index[i], style) | |
else: | |
if axis == 0: | |
for level in range(depth): | |
col += level | |
for i in range(len(index)): | |
ws.row(row + i + depth).write(col, index[i][level], style) | |
elif axis == 1: | |
for level in range(depth): | |
row += level | |
for i in range(len(index)): | |
ws.row(row).write(col + i + depth, index[i][level], style) | |
def place_data(self, ws, row=0, col=0, style=Style.default_style): | |
""" | |
Write XLtable data into an Excel sheet | |
Parameters | |
---------- | |
ws : An xlwt Worksheet object | |
row: The starting row in Excel where the data will be placed | |
(integer, base 0, default 0) | |
col: The starting column in Excel where the data will be placed | |
(integer, base 0, default 0) | |
style: An xlwt style object | |
Example | |
------- | |
>>> XLtable.place_data(ws=Sheet1, row=0, col=0, style=dstyle) | |
""" | |
for irow in range(len(self.index)): # data | |
for icol in range(len(self.columns)): | |
ws.row(row + irow).write((col + icol), | |
float(self.ix[irow][icol]), style) | |
def place_table(self, ws, row=0, col=0, rstyle=Style.default_style, | |
cstyle=Style.default_style, dstyle=Style.default_style): | |
""" | |
Write XLtable (indexes and data) into an Excel sheet | |
Parameters | |
---------- | |
ws : An xlwt Worksheet object | |
row: The starting row in Excel where the index will be placed | |
(integer, base 0, default 0) | |
col: The starting column in Excel where the index will be placed | |
(integer, base 0, default 0) | |
rstyle: An xlwt style object, determines row index style | |
cstyle: An xlwt style object, determines column index style | |
dstyle: An xlwt style object, determines data style | |
Example | |
------- | |
>>> XLtable.place_index(ws=Sheet1, row=0, col=0, rstyle=hstyle, | |
cstyle=hstyle, dstyle=data_style) | |
""" | |
self.place_index(ws=ws, row=row, col=col, axis=0, style=rstyle) | |
self.place_index(ws=ws, row=row, col=col, axis=1, style=cstyle) | |
row = row + self.row_depth | |
col = col + self.col_depth | |
self.place_data(ws=ws, row=row, col=col, style=dstyle) | |
class XLseries(pd.Series): | |
def __new__(cls, *args, **kwargs): | |
arr = pd.Series.__new__(cls, *args, **kwargs) | |
return arr.view(XLseries) | |
def __init__(self, series=None): | |
""" | |
pandas Series with extra methods for setting location and xlwt style | |
in an Excel spreadsheet. | |
Parameters | |
---------- | |
df : A pandas Series object | |
Example | |
------- | |
>>> xlvector = XLseries(series=vector) | |
""" | |
if isinstance(self.index, pd.MultiIndex): | |
self.index_depth = len(self.index.levels) | |
else: | |
self.index_depth = 1 | |
def place_index(self, ws, row=0, col=0, axis=0, style=Style.default_style): | |
""" | |
Write XLseries index into an Excel sheet | |
Parameters | |
---------- | |
ws : An xlwt Worksheet object | |
row: The starting row in Excel where the index will be placed | |
(integer, base 0, default 0) | |
col: The starting column in Excel where the index will be placed | |
(integer, base 0, default 0) | |
axis: Whether the index will be placed in vertical or horizontal | |
(0 for vertical or 1 for horizontal, default 0) | |
style: An xlwt style object | |
Example | |
------- | |
>>> XLseries.place_index(ws=Sheet1, row=0, col=0, axis=0, style=hstyle) | |
""" | |
depth = self.index_depth | |
index = self.index | |
if axis not in [0,1]: | |
raise ValueError("Excel has only two axis (0 or 1)") | |
if depth == 1: | |
if axis == 0: | |
for i in range(len(index)): | |
ws.row(row + i).write(col, index[i], style) | |
elif axis == 1: | |
for i in range(len(index)): | |
ws.row(row).write(col + i, index[i], style) | |
else: | |
if axis == 0: | |
for level in range(depth): | |
col += level | |
for i in range(len(index)): | |
ws.row(row + i).write(col, index[i][level], style) | |
elif axis == 1: | |
for level in range(depth): | |
row += level | |
for i in range(len(index)): | |
ws.row(row).write(col + i, index[i][level], style) | |
def place_data(self, ws, row=0, col=0, axis=0, style=Style.default_style): | |
""" | |
Write XLseries data into an Excel sheet | |
Parameters | |
---------- | |
ws : An xlwt Worksheet object | |
row: The starting row in Excel where the data will be placed | |
(integer, base 0, default 0) | |
col: The starting column in Excel where the data will be placed | |
(integer, base 0, default 0) | |
axis: Whether the index will be placed in vertical or horizontal | |
(0 for vertical or 1 for horizontal, default 0) | |
style: An xlwt style object | |
Example | |
------- | |
>>> XLseries.place_data(ws=Sheet1, row=0, col=0, style=dstyle) | |
""" | |
if axis == 0: | |
for i in range(len(self)): | |
ws.row(row + i).write(col, self[i], style) | |
elif axis == 1: | |
for i in range(len(self)): | |
ws.row(row).write(col + i, self[i], style) | |
def place_series(self, ws, row=0, col=0, axis=0, | |
istyle=Style.default_style, dstyle=Style.default_style): | |
""" | |
Write XLseries (index and data) into an Excel sheet | |
Parameters | |
---------- | |
ws : An xlwt Worksheet object | |
row: The starting row in Excel where the index will be placed | |
(integer, base 0, default 0) | |
col: The starting column in Excel where the index will be placed | |
(integer, base 0, default 0) | |
axis: Whether the series will be placed in vertical or horizontal | |
(0 for vertical or 1 for horizontal, default 0) | |
istyle: An xlwt style object, determines index style | |
dstyle: An xlwt style object, determines data style | |
Example | |
------- | |
>>> XLseries.place_index(ws=Sheet1, row=0, col=0, istyle=hstyle, | |
dstyle=data_style) | |
""" | |
self.place_index(ws=ws, row=row, col=col, axis=axis, style=istyle) | |
if axis == 0: | |
col = col + self.index_depth | |
else: | |
row = row + self.index_depth | |
self.place_data(ws=ws, row=row, col=col, axis=axis, style=dstyle) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment