将WebSocket数据保存至Pandas DataFrame的问题求助

我来帮你搞定这个问题！你遇到的local variable 'df' referenced before assignment错误，核心原因是回调函数on_message无法正确访问到你想传递的df变量——哪怕用了partial，如果变量的作用域没处理好，还是会出问题。下面给你两种靠谱的解决思路：

方案一：用类封装（推荐，最清晰） #

把WebSocket连接、DataFrame和回调逻辑都封装到类里，这样df作为类的实例属性，回调函数可以直接通过self.df访问，完全避开作用域问题。

示例代码：

import websocket
import pandas as pd
import json
import threading
import time

class WebSocketDataCollector:
    def __init__(self, ws_url):
        # 初始化空的DataFrame，根据你的数据结构定义列名
        self.df = pd.DataFrame(columns=["timestamp", "value", "other_field"])
        # 线程锁，避免多线程读写df时出现数据混乱
        self.lock = threading.Lock()
        # 创建WebSocket连接，绑定回调
        self.ws = websocket.WebSocketApp(
            ws_url,
            on_message=self.on_message,
            on_error=self.on_error,
            on_close=self.on_close
        )
    
    def on_message(self, ws, message):
        # 解析WebSocket消息（假设是JSON格式，根据实际情况调整）
        data = json.loads(message)
        # 加锁操作df，保证线程安全
        with self.lock:
            # 将新数据追加到DataFrame
            # 注意：频繁用append效率低，可先存列表再批量转df，这里是简化示例
            new_row = pd.DataFrame([{
                "timestamp": data.get("ts"),
                "value": data.get("val"),
                "other_field": data.get("other")
            }])
            self.df = pd.concat([self.df, new_row], ignore_index=True)
            print(f"已追加数据，当前DataFrame行数：{len(self.df)}")
    
    def on_error(self, ws, error):
        print(f"WebSocket错误: {error}")
    
    def on_close(self, ws, close_status_code, close_msg):
        print("WebSocket连接关闭")
    
    def run(self):
        # 启动WebSocket连接
        self.ws.run_forever()

# 使用示例
if __name__ == "__main__":
    collector = WebSocketDataCollector("wss://your-websocket-url.com")
    # 单独开线程运行WebSocket，不阻塞主程序调用df
    ws_thread = threading.Thread(target=collector.run)
    ws_thread.start()
    
    # 其他函数可以通过collector.df访问数据
    # 比如每隔一段时间打印最新数据
    while True:
        time.sleep(5)
        with collector.lock:
            print("当前数据预览：")
            print(collector.df.tail())

方案二：用可变对象包装df（不用类的临时方案） #

如果不想写类，可以用一个**可变对象（比如字典）**来包装df，因为可变对象是按引用传递的，回调函数里能直接修改它的内容，不会有作用域问题。

示例代码：

import websocket
import pandas as pd
import json
from functools import partial
import threading
import time

def on_message(ws, message, df_container):
    # 从容器里取出df，加锁保证线程安全
    with df_container["lock"]:
        df = df_container["data"]
        data = json.loads(message)
        new_row = pd.DataFrame([{
            "timestamp": data.get("ts"),
            "value": data.get("val")
        }])
        # 修改容器里的df
        df_container["data"] = pd.concat([df, new_row], ignore_index=True)
        print(f"已追加数据，当前行数：{len(df_container['data'])}")

def main():
    # 初始化df并放到字典容器里，同时加入线程锁
    df_container = {
        "data": pd.DataFrame(columns=["timestamp", "value"]),
        "lock": threading.Lock()
    }
    ws_url = "wss://your-websocket-url.com"
    
    # 使用partial传递容器
    ws = websocket.WebSocketApp(
        ws_url,
        on_message=partial(on_message, df_container=df_container),
        on_error=lambda ws, err: print(f"错误: {err}"),
        on_close=lambda ws, code, msg: print("连接关闭")
    )
    
    # 启动线程运行WebSocket
    threading.Thread(target=ws.run_forever).start()
    
    # 其他函数可以通过df_container["data"]访问数据
    while True:
        time.sleep(3)
        with df_container["lock"]:
            print(df_container["data"].head())

if __name__ == "__main__":
    main()

关键优化提示 #

• 提升数据追加效率：如果WebSocket消息很频繁，每次用pd.concat会拖慢速度，建议先把数据存在列表里（比如类里的self.data_list = []），每隔N条或者固定时间再一次性转成DataFrame合并到self.df里。
• 线程安全必须重视：WebSocket的回调是在单独线程里执行的，如果其他函数同时读写df，很容易出现数据混乱，一定要用threading.Lock加锁保护。
• 适配实际数据格式：一定要根据你收到的WebSocket消息格式（比如JSON、二进制、纯文本）调整解析逻辑，不然会无法转换成DataFrame的行数据。

内容的提问来源于stack exchange，提问作者swifty