diff --git a/watermark/watermark.py b/watermark/watermark.py index 780ef0a..5ee6be1 100644 --- a/watermark/watermark.py +++ b/watermark/watermark.py @@ -35,7 +35,8 @@ def watermark(author=None, email=None, github_username=None, updated=False, custom_time=None, python=False, packages=None, hostname=False, machine=False, githash=False, gitrepo=False, gitbranch=False, - watermark=False, iversions=False, watermark_self=None): + watermark=False, iversions=False, watermark_self=None, + globals_=None): '''Function to print date/time stamps and various system information. @@ -165,8 +166,16 @@ def watermark(author=None, email=None, github_username=None, if args['gitbranch']: output.append(_get_git_branch(bool(args['machine']))) if args['iversions']: - output.append(_get_all_import_versions( - watermark_self.shell.user_ns)) + if watermark_self: + ns = watermark_self.shell.user_ns + elif globals_: + ns = globals_ + else: + raise RuntimeError( + "Either `watermark_self` or `globals_` must be provided " + "to show imported package versions." + ) + output.append(_get_all_import_versions(ns)) if args['watermark']: output.append({"Watermark": __version__})