deepai / web /file.py
zkneo's picture
1.pandas-ai.done
48d8f38
Raw
History Blame Contribute Delete
680 Bytes
import pandas as pd
import io
from functools import lru_cache
from pydantic import BaseModel
import requests
class FileBean(BaseModel):
name: str
url: str
source: str
accessToken: str = None
fileId: str = None
def __hash__(self):
return hash((self.fileId))
@lru_cache
def get_dataframes(fb: FileBean):
with open(fb.url, "rb") as f:
file = f.read()
if fb.name.endswith('.xls') or fb.name.endswith('.xlsx'):
return pd.read_excel(file, sheet_name=None)
elif fb.name.endswith('.csv'):
return {fb.name: pd.read_csv(io.BytesIO(file))}
else:
raise {"msg": "File type not supported.", "success": False}