Skip to content

Instantly share code, notes, and snippets.

@dmvianna
Last active December 10, 2015 21:48
Show Gist options
  • Save dmvianna/4497752 to your computer and use it in GitHub Desktop.
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.
"""
#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