pythoncomplete in vim - hardcode factory function returns?

Question

I'm using pythoncomplete omnicompletion in vim. It works great when I instantiate classes directly, eg

import numpy as np
x = np.ndarray(l)

then x attributes complete correctly.

But I work with numpy and matplotlib so usually use factory functions ie

x = np.zeros((2,2)) 
f = plt.figure()  
ax = f.add_subplot(111)

Is there any way I can hard code the return types of these common functions so I can complete on the returned object. (ie set up some mappings that matplotlib.pyplot.figure returns matplotlib.figure.Figure, np.zeros returns np.ndarray etc.?