utils.py

import pytz
import datetime
import io
import pandas as pd


def time_in_beijing():
    '''
    return current time in Beijing as datetime object
    '''
    tz = pytz.timezone('Asia/Shanghai')
    dt = datetime.datetime.now(tz)
    return dt


def convert_string_to_datetime(date_string, time_zone="Asia/Shanghai"):
    '''
    Convert a string to a datetime object with the timezone by default,
    Shanghai
    '''
    dt = datetime.datetime.strptime(date_string, '%Y-%m-%d %H:%M:%S.%f')
    tz = pytz.timezone(time_zone)
    dt = tz.localize(dt)
    return dt


def create_stocks_entry_from_excel(byte_string):
    '''create stock entry from excel file
    Parameters
    ----------
    byte_string: bytes
        the byte string of the excel file
    Returns
    -------
    new_stock_entry: list
        [{ticker:str, shares:int, mean_price: float, date:datetime.datetime}]
        the list of stock entry
    '''
    uploaded_df = None
    with io.BytesIO(byte_string) as f:
        uploaded_df = pd.read_excel(f)

    # throw exception if doesn't have required columns
    if not set(['证券代码', '持仓数量', '平均建仓成本', 'time_stamp']).issubset(uploaded_df.columns):
        raise Exception('Missing required columns')
    # print(uploaded_df)
    # uploaded_df = pd.read_excel()
    uploaded_df.drop(columns='Unnamed: 0', inplace=True)
    # Define the regular expression pattern to match the string endings
    pattern = r'\.(sz|sh)$'
    # Define the replacement strings for each match group
    replacements = {'.sz': '.XSHE', '.sh': '.XSHG'}
    # Use the str.replace method with the pattern and replacements
    uploaded_df['证券代码'] = uploaded_df['证券代码'].str.lower()
    uploaded_df['证券代码'] = uploaded_df['证券代码'].str.replace(
        pattern, lambda m: replacements[m.group()], regex=True)
    new_stock_entry = [
        dict(ticker=ticker, shares=shares, date=time, mean_price=mean_price)
        for ticker, shares, mean_price, time in zip(
            uploaded_df['证券代码'],
            uploaded_df['持仓数量'],
            uploaded_df['平均建仓成本'],
            pd.to_datetime(uploaded_df['time_stamp']))]
    # new_profile, error = api.update_portfolio_profile(new_stock_entry)
    print(new_stock_entry)
    return new_stock_entry


def style_number(vals):
    '''color negative number as red, positive as green
    Parameters
    ----------
    vals: df columns
        the columns to be styled
    Returns
    -------
    list
        the list of style

    '''
    return ['color: red' if v < 0 else 'color: green' for v in vals]


def create_share_changes_report(df):
    '''Create a markdown report of the share changes for certain date
    Parameters
    ----------
    df: pd.DataFrame
        the dataframe of profile for a specific date
    Returns
    -------
    markdown: str
    '''

    date_str = df.date.to_list()[0].strftime('%Y-%m-%d %H:%M:%S')
    markdown = f"### {date_str}\n\n"
    markdown += 'Ticker | Display Name | Share Changes\n'
    markdown += '--- | --- | ---\n'
    for _, row in df.iterrows():
        share_changes = row['share_changes']
        # Apply green color to positive numbers and red color to negative numbers
        if share_changes > 0:
            share_changes_str = f'<span style="color:green">{share_changes}</span>'
        elif share_changes < 0:
            share_changes_str = f'<span style="color:red">{share_changes}</span>'
        else:
            share_changes_str = str(share_changes)
        markdown += '{} | {} | {}\n'.format(row['ticker'],
                                            row['display_name'], share_changes_str)
    return markdown