invest2/process_multi_ticker_data.py at master · mashinosatoshi/invest2 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import pandas as pd
from yfinance_get_info import get_info

def process_multi_ticker_dataframe(df_multi_ticker):
    """
    複数の証券コードを含むデータフレームを処理し、各証券コードにget_info関数を適用します。

    Args:
        df_multi_ticker (pd.DataFrame): 複数の証券コードを含むデータフレーム。
                                        'Code'列で証券コードが区別されている必要があります。
                                        また、get_info関数が期待する列（'Open', 'High', 'Low', 'Close', 'Volume'など）
                                        が含まれている必要があります。

    Returns:
        pd.DataFrame: 各証券コードにget_info関数が適用され、結合されたデータフレーム。
                      処理されなかった証券コードのデータは含まれません。
    """
    if df_multi_ticker.empty:
        print("入力データフレームが空です。")
        return pd.DataFrame()

    if 'Code' not in df_multi_ticker.columns:
        raise ValueError("データフレームに'Code'列が見つかりません。")

    processed_dfs = []
    # 'Code'列でグループ化し、各グループにget_info関数を適用
    for code, group_df in df_multi_ticker.groupby('Code'):
        print(f"証券コード {code} のデータを処理中...")
        # get_info関数はデータフレームのコピーを内部で行うため、ここでは直接渡す
        processed_group_df = get_info(group_df)
        if processed_group_df is not None:
            processed_dfs.append(processed_group_df)
        else:
            print(f"証券コード {code} のデータは処理されませんでした (行数が20未満など)。")

    if processed_dfs:
        # 全ての処理済みデータフレームを結合
        return pd.concat(processed_dfs).sort_values(by=['Code', 'Date']).reset_index(drop=True)
    else:
        print("どの証券コードのデータも処理されませんでした。")
        return pd.DataFrame()

if __name__ == "__main__":
    # テスト用のダミーデータフレームを作成
    # 実際のyfinanceデータに似た構造を持つようにする
    from datetime import timedelta

    # 証券コードAのデータ (十分な行数)
    dates_a = pd.to_datetime(pd.date_range(start='2023-01-01', periods=100, freq='D'))
    data_a = {
        'Date': dates_a,
        'Open': 100 + (dates_a.dayofyear % 10) * 0.5,
        'High': 105 + (dates_a.dayofyear % 10) * 0.5,
        'Low': 98 - (dates_a.dayofyear % 10) * 0.5,
        'Close': 102 + (dates_a.dayofyear % 10) * 0.5,
        'Volume': 100000 + (dates_a.dayofyear % 10) * 1000
    }
    df_a = pd.DataFrame(data_a)
    df_a['Code'] = 'STOCK_A'
    df_a = df_a.set_index('Date') # get_infoがインデックスを期待する場合に備えて

    # 証券コードBのデータ (行数が少ない、get_infoの条件を満たさない)
    dates_b = pd.to_datetime(pd.date_range(start='2023-01-01', periods=15, freq='D'))
    data_b = {
        'Date': dates_b,
        'Open': 50 + (dates_b.dayofyear % 5) * 0.2,
        'High': 52 + (dates_b.dayofyear % 5) * 0.2,
        'Low': 49 - (dates_b.dayofyear % 5) * 0.2,
        'Close': 51 + (dates_b.dayofyear % 5) * 0.2,
        'Volume': 50000 + (dates_b.dayofyear % 5) * 500
    }
    df_b = pd.DataFrame(data_b)
    df_b['Code'] = 'STOCK_B'
    df_b = df_b.set_index('Date')

    # 証券コードCのデータ (十分な行数)
    dates_c = pd.to_datetime(pd.date_range(start='2023-02-01', periods=80, freq='D'))
    data_c = {
        'Date': dates_c,
        'Open': 200 + (dates_c.dayofyear % 8) * 1.0,
        'High': 205 + (dates_c.dayofyear % 8) * 1.0,
        'Low': 198 - (dates_c.dayofyear % 8) * 1.0,
        'Close': 202 + (dates_c.dayofyear % 8) * 1.0,
        'Volume': 200000 + (dates_c.dayofyear % 8) * 2000
    }
    df_c = pd.DataFrame(data_c)
    df_c['Code'] = 'STOCK_C'
    df_c = df_c.set_index('Date')

    # 複数の証券コードを結合したデータフレーム
    combined_df = pd.concat([df_a, df_b, df_c])
    print("結合されたダミーデータフレームのヘッド:")
    print(combined_df.head())
    print("\n結合されたダミーデータフレームのテール:")
    print(combined_df.tail())
    print(f"\n結合されたダミーデータフレームの形状: {combined_df.shape}")
    print(f"ユニークな証券コード: {combined_df['Code'].unique()}")

    # 関数を呼び出して処理
    print("\n--- process_multi_ticker_dataframe を実行中 ---")
    processed_combined_df = process_multi_ticker_dataframe(combined_df)

    print("\n--- 処理結果 ---")
    if not processed_combined_df.empty:
        print("処理済みデータフレームのヘッド:")
        print(processed_combined_df.head())
        print("\n処理済みデータフレームのテール:")
        print(processed_combined_df.tail())
        print(f"\n処理済みデータフレームの形状: {processed_combined_df.shape}")
        print(f"処理されたユニークな証券コード: {processed_combined_df['Code'].unique()}")
        # 特定の証券コードのデータを確認
        print("\nSTOCK_A の処理済みデータ (一部):")
        print(processed_combined_df[processed_combined_df['Code'] == 'STOCK_A'].head())
        print("\nSTOCK_C の処理済みデータ (一部):")
        print(processed_combined_df[processed_combined_df['Code'] == 'STOCK_C'].head())
    else:
        print("処理結果のデータフレームは空です。")

    # 空のデータフレームを渡した場合のテスト
    print("\n--- 空のデータフレームでテスト ---")
    empty_df = pd.DataFrame(columns=['Date', 'Open', 'High', 'Low', 'Close', 'Volume', 'Code'])
    processed_empty_df = process_multi_ticker_dataframe(empty_df)
    print(f"空のデータフレーム処理結果の形状: {processed_empty_df.shape}")

    # 'Code'列がないデータフレームを渡した場合のテスト
    print("\n--- 'Code'列がないデータフレームでテスト ---")
    df_no_code = pd.DataFrame({'Date': dates_a, 'Close': 100 + (dates_a.dayofyear % 10) * 0.5})
    try:
        process_multi_ticker_dataframe(df_no_code)
    except ValueError as e:
        print(f"エラーを捕捉しました: {e}")