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是 %%