从另一个模块注册jupyter变量

Question

我创建了一个简单的“魔法”函数，用来调用AWS Athena。

import awswrangler as wr
import inspect
import pandas as pd
from IPython.core.magic import (register_line_magic, register_cell_magic,
                                register_line_cell_magic)


@register_cell_magic
def athena(line, cell):
    """
    Execute athena query and return the result as a pandas dataframe

    Usage:
    >>> %%athena <database> [<destination_var>]

    :param line:
    :param cell:
    :return:
    """

    args = []
    if line:
        args = line.split()
    df = wr.athena.read_sql_query(cell, database=args[0])
    if len(args) > 1:
        globals()[args[1]] = df
    else:
        return df

然后我把它放到了一个单独的模块里，这个模块是从Jupyter笔记本中导入的。

问题是，生成的数据表（dataframe）并没有被注册到Jupyter笔记本的全局变量中，而是注册到了那个模块的全局变量里。

当我把这段代码直接写在笔记本里并执行时，变量就能正确注册并且可以访问了。

那么，怎样才能把这个数据表正确地注册到Jupyter的全局变量中呢？

我尝试了几种方法，使用inspect库来获取导入模块，但都没有成功，我觉得这不是个有效的解决办法。

难道就没有什么办法可以正确获取笔记本的全局变量，或者有没有什么特殊的上下文可以让我正确注册变量呢？

相关问题：

• 确认 import * 和 from xxx import * 之间的区别
• 在异常时启动IPython shell
• 在Python中，如何从导入模块访问主模块的命名空间？
• 全局作用域中的描述符？
• 使用上下文导入模块
• 来自其他文件的全局变量 Python
• 动态加载模块后对加载模块使用"from x import *
• Python - 将全局变量赋值为函数返回值需要函数也全局吗？

Answer 1

在阅读了这个链接的内容后，我找到了答案，所以正确的实现方式是：

import awswrangler as wr
from IPython.core.magic import (register_cell_magic,
                                needs_local_scope)


@register_cell_magic
@needs_local_scope
def athena(line, cell, local_ns=None):
    """
    Execute athena query and return the result as a pandas dataframe

    Usage:
    >>> %%athena <database> [<destination_var>]

    :param line:
    :param cell:
    :return:
    """

    print(local_ns.keys())

    args = []
    if line:
        args = line.split()
    df = wr.athena.read_sql_query(cell, database=args[0])
    if len(args) > 1:
        local_ns[args[1]] = df
    else:
        return df