Jupyter生態二次開發系列(五)

今天記錄一下如何開發jupyter自定義magic, 就是jupyter裏面%%或者!這樣用來標識解釋器環境的東西.

如下代碼以訪問 greenplum 數據庫爲例.

import psycopg2
import psycopg2.extras
from configparser import ConfigParser
import pandas as pd
from IPython.core.magic import (
    cell_magic,
    Magics,
    magics_class
)
import os

# Ipython的修飾方法修飾該class是magic
@magics_class
class GreenPlumMagix(Magics):
    # 定義默認鏈接參數
    pg_default = {
        'dbname': 'dmp',
        'host': 'dmp-gp1',
        'port': '5432',
        'user': 'guest',
        'password': 'xxxxxxx'
    }

    def connect(self):
        # 如下路徑爲自定義magic及配置文件的存放標準路徑
        conf_file = os.environ['HOME'] + '/.ipython/profile_default/startup/postgresql_magic.ini'
        cp = ConfigParser()
        if os.path.exists(conf_file):
            #若是存在ini配置文件則從配置文件中讀取greenplum配置, 不然建立該ini文件
            try:
                cp.read(conf_file)
                if 'postgresql' in cp.sections():
                    return psycopg2.connect(
                        dbname=cp['postgresql']['dbname'],
                        host=cp['postgresql']['host'],
                        port=int(cp['postgresql']['port']),
                        user=cp['postgresql']['user'],
                        password=cp['postgresql']['password']
                    )
                else:
                    cp['postgresql'] = self.pg_default
                    with open(conf_file, 'a+') as f:
                        cp.write(f)
                    return psycopg2.connect(
                        dbname=self.pg_default['dbname'],
                        host=self.pg_default['host'],
                        port=int(self.pg_default['port']),
                        user=self.pg_default['user'],
                        password=self.pg_default['password']
                    )
            except Exception as e:
                print(e)
        else:
            cp['postgresql'] = self.pg_default
            with open(conf_file, 'a+') as f:
                cp.write(f)
            # 返回greenplum鏈接對象
            return psycopg2.connect(
                dbname=self.pg_default['dbname'],
                host=self.pg_default['host'],
                port=int(self.pg_default['port']),
                user=self.pg_default['user'],
                password=self.pg_default['password']
            )

    使用 cell_magic 修飾greenplum方法, 還能夠用 line_magic或者line_cell_magic
    @cell_magic
    def greenplum(self, line='', cell=None):
        sql = cell
        k = self.connect()
        kursor = k.cursor()
        kursor.execute(sql)
        # 將返回數據放入 pandas, 這樣在jupyter頁面返回時會更美觀一些
        try:
            pd.set_option('display.max_columns', None)
            pd.set_option('display.max_rows', None)
            results = kursor.fetchall()
            return pd.DataFrame(results)
        except:
            return pd.DataFrame([])

# 註冊自定義magic類到ipython裏面
ipy = get_ipython()
ipy.register_magics(GreenPlumMagix)

再打開notebook就能夠了

使用方式以下, line magic是 %, cell magic是 %%