Python覚え書き

久しぶりに結構がっつり書いたら色々と新しい知見があり、コードを見れば分かると言えば分かるわけだが、まあ次に書く時には忘れちゃうので、まとめておく。これが最良・最新という訳では決してないと思うのだけれど、とりあえず目的を果たすものはできた、と。

argparse

argparseを使ったサブコマンド

ツールというかコマンドをPythonで実装する際に、argparseを必ず使うのだけれど、久々だったのでサブコマンドの書き方を忘れていた。

import argparse
import sys
from functools import wraps
from typing import Callable

def task_a(subparser) -> None:
    pass

def map_subparser_to_func(func, subparser) -> Callable:
    @wraps(func)
    def wrapper(*args, **kwargs):
        return func(subparser, *args, **kwargs)
    return wrapper

def main() -> None:
    parser = argparse.ArgumentParser(description = "description of the command")
    # Global argument
    parser.add_argument()

    subparsers = parser.add_subparsers(dest = "subparser")

    parser_a = subparsers.add_parser("task_a")
    parser_a.add_argument(...)
    parser_a.set_defaults(func = map_subparser_to_func(task_a, parser_a)

    args = parser.parse_args(sys.argv[1:])

    # subcommand is required
    if args.subparser != None:
        args.func()
    else:
        parser.print_help()

if __name__ == "__main__":
    main()

これで、サブコマンド”task_a”は、task_a()に割り当てられる。

psycopg

psycopg2ではなくpsycopg

ながらくpsycopg2を使っていたが、ひょっとしてバージョン3があるんじゃないかと思って調べたら、バージョン3はバージョン番号無しになってる。

# Old
import psycopg2

# New
import psycopg

psycopg2のDictCursorの代替

パフォーマンスが良くなっている一方、psycopg2で使えるDictCursorは無くなってた。代わりに、connect()もしくはcursor()で指定する。下記の例だと、row_factoryで指定している。

import sys

try:
    import psycopg as pg
    from psycopg.rows import dict_row, namedtuple_row
except:
    cur_python = f"python{sys.version_info[0]}.{sys.version_info[1]}"
    sys.exit(f'''Please install 'psycopg' before executing.\n
        {cur_python} -m pip install psycopg''')

def connect_db(connect_string: str) -> pg.connection:
    try:
        print("Connecting database...")
        connection = pg.connect(connect_string)
        return connection
    except:
        exit_with_error('''Failed to connect to database.\n
            Please check connect string you passed.''')

connect_string = "host=localhost dbname=postgres"

with connect_db(connect_string) as con:
    with con.cursor(row_factory = namedtuple_row) as cur:
        cur.execute("SELECT foo, bar FROM tables;")
        rows = cur.fetchall()
        if len(rows) > 0:
            for row in rows:
                prinf(row.foo)
                prinf(row.bar)

リストからPostgreSQLの配列型に

リストをPostgreSQLの配列型として、文字列でINSERT / UPDATEしたい時がある。

def list_to_pgsql_array(ls) -> str:
    if type(ls) is str:
        ls = [ls]

    return "{" + ",".join(map(lambda s: '"' + s + '"', sorted(set(ls), key = ls.index))) + "}"

openpyxl

openpyxlの警告の抑制

openpyxlでExcelファイルを読み書きする際に、”Workbook contains no default style, apply openpyxl’s default”という警告が出てうるさい。下記のように抑制する。

import warnings

# suppress warnings from openpyxl
warnings.filterwarnings('ignore', category=UserWarning, module='openpyxl')

ハイパーリンク(macOS限定?)

セルにハイパーリンクを設定する際、下記のコードなら動くが、他のコードだとうまくいかない。

from openpyxl.styles import Font

wb = xl.Workbook()
ws = wb.create_sheet("Sheet 1", 0)

ws.cell(row = 1, column = 1).value = "リンク"
ws.cell(row = 1, column = 1).hyperlink = "https://example.com/"
ws.cell(row = 1, column = 1).font = Font(u = "single", color = "000563C1")

シート全体のフォントの指定

valueを設定したセル全てに対してフォントを適用するコードは見かけるが、個別のセルにフォント書式を設定して最後にセル全てにフォントを適用すると、全て書式が消えてしまう。ワークシートのデフォルトフォントを設定するコードがなかなか分からなかった。以下のハックが有効。

from openpyxl.styles import Font, DEFAULT_FONT

wb = xl.Workbook()
ws = wb.create_sheet("Sheet 1", 0)

# Filling cells here

_font = Font(name="游ゴシック")
{k: setattr(DEFAULT_FONT, k, v) for k, v in _font.__dict__.items()}

wb.save("~/")

行をなめる高速な方法

ワークシートを上から順になめて行くときに、イテレータを使わないと遅い。

import openpyxl as xl

wb = xl.load_workbook(workbook_path)
ws = wb["Sheet 1"]
ws = wb.active

for r in range(1, ws.max_row):
    value = ws.cell(row = r, column = 1).value

以下のように書くと、高速。

for c in ws.iter_rows(min_row = 2, values_only = True):
    value = c[0]

Tableの設定

Excel内でのテーブルの設定。

table = Table(displayName = "Table1", ref = "A1:M100")
table.tableStyleInfo = TableStyleInfo(name = "TableStyleMedium2", showRowStripes = True)
ws.add_table(table)

subprocess

これも忘れてた。subprocessによるコマンドの実行結果のリダイレクション。

with open("result.txt", "w") as fo:
    subprocess.run(['ls', '-1'], stdout = fo)

その他

毎回忘れて検索するはめになる、ターミナルでの書式設定。以下で赤文字

print("\033[31m" + "a message" + "\033[0m")