From adcdea85fd3c2cba3016342caf1cc6c88a0dc915 Mon Sep 17 00:00:00 2001
From: Gus Gordon <gusgordon@users.noreply.github.com>
Date: Wed, 14 Dec 2016 12:57:53 -0500
Subject: [PATCH] ENH: Format zipline assets before displaying (#358)

* Add new function for formatting zipline assets. Apply it whenever asset names are displayed.

* Check for general asset

* Update doc

* Use latest PyMC3 for tests
---
 .travis.yml            |  3 ++-
 pyfolio/plotting.py    | 10 ++++++----
 pyfolio/round_trips.py |  3 ++-
 pyfolio/utils.py       | 18 ++++++++++++++++++
 4 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 7216cd57..96ca5bc1 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -28,7 +28,8 @@ install:
   - source activate testenv
   - pip install nose_parameterized
   #- pip install --no-deps git+https://github.com/quantopian/zipline
-  - pip install -e .[bayesian]
+  - pip install -e .
+  - pip install git+git://github.com/pymc-devs/pymc3.git
 
 before_script:
   - "flake8 pyfolio"
diff --git a/pyfolio/plotting.py b/pyfolio/plotting.py
index 7959e430..bccdff76 100644
--- a/pyfolio/plotting.py
+++ b/pyfolio/plotting.py
@@ -1048,6 +1048,8 @@ def show_and_plot_top_positions(returns, positions_alloc,
         The axes that were plotted on.
 
     """
+    positions_alloc = positions_alloc.copy()
+    positions_alloc.columns = positions_alloc.columns.map(utils.format_asset)
 
     df_top_long, df_top_short, df_top_abs = pos.get_top_long_short_abs(
         positions_alloc)
@@ -1733,7 +1735,7 @@ def plot_round_trip_lifetimes(round_trips, disp_amount=16, lsize=18, ax=None):
                     linewidth=lsize, solid_capstyle='butt')
 
     ax.set_yticks(range(disp_amount))
-    ax.set_yticklabels(sample)
+    ax.set_yticklabels([utils.format_asset(s) for s in sample])
 
     ax.set_ylim((-0.5, min(len(sample), disp_amount) - 0.5))
     blue = patches.Rectangle([0, 0], 1, 1, color='b', label='Long')
@@ -1765,10 +1767,10 @@ def show_profit_attribution(round_trips):
     """
 
     total_pnl = round_trips['pnl'].sum()
-    pct_profit_attribution = round_trips.groupby(
-        'symbol')['pnl'].sum() / total_pnl
+    pnl_attribution = round_trips.groupby('symbol')['pnl'].sum() / total_pnl
 
-    utils.print_table(pct_profit_attribution.sort_values(
+    pnl_attribution.index = pnl_attribution.index.map(utils.format_asset)
+    utils.print_table(pnl_attribution.sort_values(
         inplace=False,
         ascending=False),
         name='Profitability (PnL / PnL total) per name',
diff --git a/pyfolio/round_trips.py b/pyfolio/round_trips.py
index 5fce367a..1a2e7a32 100644
--- a/pyfolio/round_trips.py
+++ b/pyfolio/round_trips.py
@@ -20,7 +20,7 @@
 import pandas as pd
 import numpy as np
 
-from .utils import print_table
+from .utils import print_table, format_asset
 
 PNL_STATS = OrderedDict(
     [('Total profit', lambda x: x.sum()),
@@ -404,5 +404,6 @@ def print_round_trip_stats(round_trips, hide_pos=False):
                 name='Return stats')
 
     if not hide_pos:
+        stats['symbols'].columns = stats['symbols'].columns.map(format_asset)
         print_table(stats['symbols'] * 100,
                     fmt='{:.2f}%', name='Symbol stats')
diff --git a/pyfolio/utils.py b/pyfolio/utils.py
index 5db2d7af..1d9c6c6b 100644
--- a/pyfolio/utils.py
+++ b/pyfolio/utils.py
@@ -133,6 +133,24 @@ def _1_bday_ago():
     return pd.Timestamp.now().normalize() - _1_bday
 
 
+def format_asset(asset):
+    """
+    If zipline asset objects are used, we want to print them out prettily
+    within the tear sheet. This function should only be applied directly
+    before displaying.
+    """
+
+    try:
+        import zipline.assets
+    except:
+        return asset
+
+    if isinstance(asset, zipline.assets.Asset):
+        return asset.symbol
+    else:
+        return asset
+
+
 def get_returns_cached(filepath, update_func, latest_dt, **kwargs):
     """
     Get returns from a cached file if the cache is recent enough,