【Python】in 検索した結果のカラムのみ表示させる

経緯

先日試しで作った関数が結構便利かなっと思いましたのでメモ ✍️
250カラムほどある中から、"賞味期限xx"のようなカラムだけ抽出した

やりたいこと

• df[["賞味期限1", "賞味期限2",:...]] のような手打ちが減る 👈 ここを改善したかった
• "賞味期限1", "賞味期限2" :... のようなカラムがある場合一度に検索抽出できる

作成した関数

# in 検索した結果のカラムのみ表示させる

def in_search_col_list(df, cols="賞味期限", head=2) -> List:
    """ in search した colの結果を返す

    params
    ----------
    df(DataFrame): 取り込みデータ
    cols(str): 検索にかけたいグローバルの値
    head(int): 先頭から抽出するレコード数

    return
    ----------
    df(IndexList) in search したカラムリスト結果を返す
    """
    return df[[for col in list(df.columns) if cols in col]].head(head)

使用例

# 必要package
import pandas as pd
import numpy as np
from typing import List

# サンプルデータ（実際は, *tabファイルを読み込んでいたりします。）
df = pd.DataFrame({ '賞味期限A' : 1.,
                    '保管場所B' : pd.Timestamp('20130102'),
                    '保管場所C' : pd.Series(1,index=list(range(5)),dtype='float32'),
                    '賞味期限D' : np.array([3] * 5,dtype='int32'),
                    'E賞味期限' : pd.Categorical(["test","train","test","train", "test"]),
                    '食品個F' : 'foo',
                    '期間限定G' : ['1000', np.nan, '123', '1234', '']})


def in_search_col_list(df, cols="賞味期限", head=2) -> List:
    return df[[ col for col in list(df.columns) if cols in col]].head(head)

# 関数呼び出し
in_search_col_list(df, cols="賞味期限", head=2)

出力結果

まとめ

• cols の引数名に検索したい文字列を入れれば、一発で検索できるので便利