Allow code parameter to be a function

[jmanj78]

Currently the code parameter must be a static string, it would be helpful if it could instead be a function with the same signature as the Storybook render function