-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathexample3.py
117 lines (91 loc) · 2.92 KB
/
example3.py
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
# NOTE: WORK IN PROGRESS
from datetime import date
from enum import Enum
# from typing import Dict
# from typing import Tuple
from typing import Optional
from typing import Set
import altair as alt
# import numpy as np
import pandas as pd
import streamlit as st
from pydantic import BaseModel
from pydantic import Field
from pydantic import validator
from statelit import StateManager
from statelit.types import DateRange
df = pd.read_csv("assets/stock_prices.csv")
df["date"] = pd.to_datetime(df["date"]).dt.date
print(df["date"].iloc[0], type(df["date"].iloc[0]))
# df["log_price"] = np.log(df["price"])
stocks_list = df["ticker"].unique()
Stock = Enum("Stock", {i: i for i in stocks_list}, type=str)
class PriceType(str, Enum):
OPEN = "open"
CLOSE = "close"
HIGH = "high"
LOW = "low"
ADJ_CLOSE = "adj_close"
class DashboardConfig(BaseModel):
class Config:
@staticmethod
def streamlit_label_generator(x):
return x.replace("_", " ").title()
# TODO: manually editing Tuple[date, date] doesn't work
date_range: DateRange = DateRange(date(2021, 1, 1), date(2022, 12, 31))
price_type: PriceType = "close"
stocks: Set[Stock] = ["AAPL"]
log_scale: bool = False
# TODO: When validating, error message happens twice
index_to_start_date: Optional[date] = Field(default=date(2021, 3, 2), streamlit_disabled=True)
@validator("index_to_start_date", allow_reuse=True)
def validate_trading_date(cls, v):
if v is None:
return None
if v not in df["date"].to_numpy():
raise ValueError("Not a valid trading date")
return v
state_manager = StateManager(DashboardConfig)
with st.sidebar:
st.subheader("Dashboard config")
config = state_manager.form()
st.markdown("---")
state_manager.text_area()
st.text("\0\n\n\0") # padding
if config.index_to_start_date:
_df = df.loc[
df["date"] == config.index_to_start_date,
["ticker", "price_type", "price"]
].rename(columns={"price": "index_base_price"})
df = df.merge(
right=_df,
on=["ticker", "price_type"],
how="left"
)
df["price_index"] = df["price"] / df["index_base_price"]
df = df.loc[
(df["price_type"] == config.price_type)
& (df["ticker"].isin(config.stocks))
& (df["date"] >= config.date_range[0])
& (df["date"] <= config.date_range[1])
]
col_name = "price" if not config.index_to_start_date else "price_index"
chart = (
alt.Chart(df)
.mark_line()
.encode(
x="date:T",
y=alt.Y(col_name, scale=alt.Scale(type="log" if config.log_scale else "linear")),
color="ticker"
)
)
if config.index_to_start_date:
vertical_line_for_date = (
alt.Chart(pd.DataFrame([{"date": config.index_to_start_date}]))
.mark_rule()
.encode(
x="date"
)
)
chart += vertical_line_for_date
st.altair_chart(chart, use_container_width=True)