Setting axes height - stretch axes height to legend height - python-2.7

Question

I have a plot with a legend whose height is bigger than the axes height (like the result of the code below). Now, I would like to stretch the axes height in such a way that it ends at the legend's end.
import matplotlib.pyplot as plt
import numpy as np
t = np.arange(0., 10.5, 0.5)
fig, ax = plt.subplots()
for c in range(0, 20):
ax.plot(t, t*c/2, label='func {}'.format(c))
ax.legend(bbox_to_anchor=(1.01, 1), loc=2, borderaxespad=0.)
What I would like to have is something like that (it does not matter, if the bottom of the grid or the bottom of the ticks ends with the legend)
I tried to reach my goal with appending following code but without any effect (legend_h is always =1.0):
legend_h = ax.get_legend().get_window_extent().height
ax_height = ax.get_window_extent().height
if ax_height < legend_h:
fig.set_figheight(legend_h/ax_height * fig.get_figheight())
Further on, it would be nice if I could only change the properties of the axes itself and not of the whole figure.
Edit:
My main purpose is to run the figure generation from a script but I also tried it in the Ipython notebook. One try was also to temporarily store the figure before getting the heights and setting the new figure height. But that also did not produce correct results.

Answer 1

I think you can achieve what you want by simply adding plt.draw() to what you already have, e.g.
fig, ax = plt.subplots()
for c in range(0, 20):
ax.plot(t, t*c/2, label='func {}'.format(c))
ax.legend(bbox_to_anchor=(1.01, 1), loc=2, borderaxespad=0.)
plt.draw()
legend_h = ax.get_legend().get_window_extent().height
ax_height = ax.get_window_extent().height
if ax_height < legend_h:
fig.set_figheight(legend_h/ax_height * fig.get_figheight())
Update: Also, you can try (which should work from a script, and based on this answer):
import matplotlib.pyplot as plt
import numpy as np
t = np.arange(0., 10.5, 0.5)
fig, ax = plt.subplots()
for c in range(0, 20):
ax.plot(t, t*c/2, label='func {}'.format(c))
lgd = ax.legend(bbox_to_anchor=(1.01, 1), loc=2, borderaxespad=0.)
fig.tight_layout()
fig.savefig('script.png', bbox_extra_artists=(lgd,), bbox_inches='tight')

Answer 2

In principle, #Matt Pitkin's answer shows the right approach. However, rather than set_figheight one would use set_size_inches. The calculation also needs to include the figure margins, which can be obtained from the fig.subplotpars.
Additionally to the height, we can also set the width of the figure, such that the legend is included.
import matplotlib.pyplot as plt
import numpy as np
t = np.linspace(0,10); c=20
fig, ax = plt.subplots()
for c in range(0, 20):
ax.plot(t, t*c/2, label='func {}'.format(c))
bbox = (1.01,1)
ax.legend(bbox_to_anchor=bbox, loc=2, borderaxespad=0.)
fig.canvas.draw()
legend_h = ax.get_legend().get_window_extent().height
ax_height = ax.get_window_extent().height
if ax_height < legend_h:
w,h = fig.get_size_inches()
h =legend_h/fig.dpi/(fig.subplotpars.top-fig.subplotpars.bottom)
fig.set_size_inches(w,h)
# set width as well
w,h = fig.get_size_inches()
r = ax.get_legend().get_window_extent().width/fig.dpi/w
fig.subplots_adjust(right=1-1.1*r)
plt.show()
The picture below is when running this as a script.
In Ipython or jupyter, the figure will automatically be cropped or expanded, because the png shown is automatically saved using the bbox_inches='tight' option. Therefore, the width adjustment is not necessary for a jupyter notebook.