Python實現(xiàn)自定義Jupyter魔法命令

更新時間：2022年08月19日 09:47:33 作者：古明地覺

相信大家都用過?jupyter，也用過里面的魔法命令，這些魔法命令都以%或者%%開頭。用法還是比較簡單的，但是我們能不能自定義魔法命令呢？本文就來教大家如何自定義Jupyter魔法命令

相信大家都用過 jupyter，也用過里面的魔法命令，這些魔法命令都以 % 或者 %% 開頭，我們舉個例子。

用法還是比較簡單的，但是我們能不能自定義魔法命令呢？毫無疑問是可以的，因為上面的 %%cython 就是 Cython 模塊自定義的。

所以命令可以是 jupyter 內(nèi)置的，比如 %time，直接拿來就能用；還可以是第三方模塊里面的，在 jupyter 通過 %load_ext 加載之后，再嵌入進來。下面就來看看如何自定義魔法命令。

from?IPython.core.magic?import?(
????magics_class,
????Magics,
????line_magic,
????cell_magic
)


@magics_class
class?MagicOrder(Magics):
????"""
????自定義一個類，類名叫什么無所謂
????但要繼承?Magics，并且要被?magics_class?裝飾
????"""

????@line_magic
????def?hello(self,?line):
????????"""
????????在 jupyter 中就可以使用如下命令，比如：
????????%hello?<Your?Code>，然后就會調(diào)用這個?hello?方法
????????參數(shù)?line?就是?%hello?后面的代碼
????????"""
????????print(f"line:?{line}")

????@cell_magic
????def?world(self,?line,?cell):
????????"""
????????在 jupyter 中就可以使用如下命令，比如：
????????%%world
????????<Your?Code>
????????<Your?Code>
????????...

????????然后就會調(diào)用這個?world?方法
????????參數(shù)?cell?就是?%%world?下面整個單元格的代碼

????????然后還有一個參數(shù)?line，它表示?%%world?所在行后面的代碼
????????但對于?%%?開頭的命令來說，我們一般都會新起一行，然后寫代碼
????????所以?line?這個參數(shù)暫時用不到
????????"""
????????print(f"line:?\n{line}")
????????print("-----------------")
????????print(f"cell:?\n{cell}")


#?必須定義?load_ipython_extension?函數(shù)
#?%load_ext?本質(zhì)上也是加載一個模塊
#?但它會自動調(diào)用該函數(shù)
def?load_ipython_extension(ip):
????#?在函數(shù)內(nèi)部，我們將類?MagicOrder?注冊進去
????#?然后就可以使用它內(nèi)部的魔法命令了
????ip.register_magics(MagicOrder)

#?如果不定義此函數(shù)，那么使用?%load_ext?加載時會報錯
#?The?xxx?module?is?not?an?IPython?extension.

當(dāng)前模塊叫 main.py，我們來測試一下：

結(jié)果沒有問題，但說實話對于 %% 開頭的命令來說，我們很少會在它后面寫代碼，基本都是新起一行，就像下面這個樣子。

自定義命令我們已經(jīng)實現(xiàn)了，并且也知道怎么獲取輸入的代碼了，下面要做的就是執(zhí)行它。而將字符串當(dāng)成代碼執(zhí)行，我們可以使用內(nèi)置函數(shù) exec。

@magics_class
class?MagicOrder(Magics):

????@line_magic
????def?hello(self,?line):
????????exec(line)

????@cell_magic
????def?world(self,?line,?cell):
????????exec(cell)

代碼的其它部分不變，然后你覺得接下來調(diào)用魔法命令會執(zhí)行成功嗎？我們測試一下：

神奇的地方出現(xiàn)了，雖然命令執(zhí)行成功了，但執(zhí)行完之后，告訴我們變量未定義。其實原因很好想，我們調(diào)用 exec 的時候沒有指定名字空間，那么默認會影響 exec 函數(shù)所在的名字空間，即 hello 和 world 函數(shù)的名字空間。

當(dāng)打開一個 jupyter 的時候，內(nèi)部相當(dāng)于啟動了一個 shell，所以在調(diào)用 exec 的時候，應(yīng)該將整個 shell 的名字空間傳進去。

from?IPython.core.magic?import?(
????magics_class,
????Magics,
????line_magic,
????cell_magic,
????needs_local_scope
)


@magics_class
class?MagicOrder(Magics):

????@line_magic
????def?hello(self,?line):
????????#?通過?self.shell.user_ns
????????#?可以拿到當(dāng)前?shell?的名字空間
????????#?注意：包含所有的單元格
????????local_ns?=?self.shell.user_ns
????????#?在?local_ns?當(dāng)中執(zhí)行代碼
????????exec(line,?local_ns,?local_ns)

????@needs_local_scope
????@cell_magic
????def?world(self,?line,?cell,?local_ns):
????????#?或者通過?needs_local_scope?裝飾器
????????#?這樣在調(diào)用函數(shù)的時候，會額外傳遞一個?local_ns?參數(shù)
????????#?該參數(shù)和?self.shell.user_ns?等價
????????exec(cell,?local_ns,?local_ns)

def?load_ipython_extension(ip):
????ip.register_magics(MagicOrder)

然后再來測試一下：

此時就沒有任何問題了。

下面我們模仿 jupyter 的 %time 命令，實現(xiàn)一個 %my_time，來加深一遍印象。

@magics_class
class?MagicOrder(Magics):

????@needs_local_scope
????@line_magic
????def?my_time(self,?line,?local_ns):
????????start?=?time.perf_counter()
????????exec(line,?local_ns,?local_ns)
????????end?=?time.perf_counter()
????????print(f"總耗時:?{round(end?-?start,?3)}")

測試一下：

結(jié)果沒有問題，是我們想要的結(jié)果。

最后再來看看如何設(shè)置可選參數(shù)，舉一個 Cython 的例子：

我們說對于以 %% 開頭的命令，應(yīng)該新起一行，在它的下面寫代碼。而之所以新起一行，是因為命令所在的行，要用于設(shè)置可選參數(shù)。那么問題來了，如何設(shè)置指定的可選參數(shù)呢？

from?IPython.core.magic?import?(
????magics_class,
????Magics,
????cell_magic,
????needs_local_scope
)
from?IPython.core?import?magic_arguments


@magics_class
class?MagicOrder(Magics):
????@magic_arguments.magic_arguments()
????#?在?jupyter?中可以通過?-n=xxx?或者?--name=xxx
????#?然后是?dest="name"，用于指定參數(shù)的名字
????#?后續(xù)便可以通過?name?字段來獲取該參數(shù)的值
????@magic_arguments.argument(
????????"-n",?"--name",?dest="name",?default="satori"
????)
????#?"-"?和?"--"?可以只出現(xiàn)一個，并且默認解析得到的是字符串
????#?而?age?我們希望是整數(shù)，所以指定?type?為?int
????#?解析完參數(shù)之后，會自動調(diào)用?int?進行轉(zhuǎn)化
????#?如果不指定該參數(shù)，則使用?default
????#?而這里沒有?default，那么結(jié)果就是?None
????@magic_arguments.argument(
????????"--age",?dest="age",?type=int
????)
????@magic_arguments.argument(
????????"-h",?"--hobby",?dest="hobby",?default=[],
????????action="append"
????)
????@needs_local_scope
????@cell_magic
????def?order(self,?line,?cell,?local_ns):
????????#?顯然?line?就是可選參數(shù)，cell?就是代碼塊
????????exec(cell,?local_ns,?local_ns)
????????#?解析參數(shù)
????????args?=?magic_arguments.parse_argstring(
????????????self.order,?line)
????????#?打印
????????print(args)


def?load_ipython_extension(ip):
????ip.register_magics(MagicOrder)

我們測試一下：