Files
stockdb/fs_datasets.py
2024-09-10 22:23:14 -04:00

63 lines
2.1 KiB
Python

import requests
import zipfile
import os
import pandas as pd
def download_file(url, filename):
headers = {"User-Agent": "LeonardExcoffier/1.0 (excoffier.leonard@gmail.com)"}
response = requests.get(url, headers=headers, stream=True)
response.raise_for_status()
with open(filename, "wb") as file:
for chunk in response.iter_content(chunk_size=8192):
file.write(chunk)
def extract_zip(source_filename, destination_folder):
with zipfile.ZipFile(source_filename, "r") as zip_ref:
zip_ref.extractall(destination_folder)
dataset_dir = "dataset"
dataset_link = "https://www.sec.gov/files/dera/data/financial-statement-data-sets/2024q2.zip"
dataset_zip = "2024q2.zip"
#os.makedirs(dataset_dir, exist_ok=True)
#download_file(dataset_link, dataset_zip)
#extract_zip(dataset_zip, dataset_dir)
num_file_path = os.path.join(dataset_dir, "num.txt")
pre_file_path = os.path.join(dataset_dir, "pre.txt")
sub_file_path = os.path.join(dataset_dir, "sub.txt")
tag_file_path = os.path.join(dataset_dir, "tag.txt")
# List of file paths and corresponding DataFrames
file_paths = {
"sub": os.path.join(dataset_dir, "sub.txt"),
"tag": os.path.join(dataset_dir, "tag.txt"),
"num": os.path.join(dataset_dir, "num.txt"),
"pre": os.path.join(dataset_dir, "pre.txt"),
}
# Dictionary to hold DataFrames
dfs = {}
# Primary keys for each DataFrame
primary_keys = {
"sub": ["adsh"],
"tag": ["tag", "version"],
"num": ["adsh", "tag", "version", "ddate", "qtrs", "uom"],
"pre": ["adsh", "report", "line"],
}
# Loop to read each file and create a DataFrame
for key, path in file_paths.items():
dfs[key] = pd.read_csv(path, sep='\t')
# Drop rows where primary key columns are empty
dfs[key].dropna(subset=primary_keys[key], inplace=True)
# Drop duplicate rows based on primary key columns
dfs[key].drop_duplicates(subset=primary_keys[key], inplace=True)
print(f"{key.upper()} DataFrame after dropping empty primary keys and duplicates:")
print(dfs[key].info())
print("\n")