refactor add_details_to_stock_df to unitls.py

api.py

@@ -10,6 +10,8 @@ from typing import List, Optional
 from sqlalchemy import create_engine
 import table_schema as ts
 import os
+import utils
+
 db_url = 'sqlite:///instance/local.db'
 load_dotenv()
 user_name = os.environ.get('JQDATA_USER')
@@ -266,6 +268,7 @@ def fetch_stocks_price(**params):
     return stocks_df
 
 
+
 @auth_api
 def update_benchmark_profile(start_date: datetime,
                              end_date: datetime,
@@ -338,9 +341,9 @@ def update_benchmark_profile(start_date: datetime,
         subset=['ticker', 'date'], keep='last', inplace=True)
 
     # update deail
-    incoming_error = add_detail_to_stocks(update_df)
-    error_message.extend(incoming_error)
+    merged_df = utils.add_details_to_stock_df(update_df)
+    # error_message.extend(incoming_error)
 
-    return update_df, error_message
+    return merged_df, error_message
 
 # get_all_stocks_detail()

utils.py

@@ -2,16 +2,45 @@ import pytz
 import datetime
 import io
 import pandas as pd
+import table_schema as ts
+from sqlalchemy import create_engine
 
+db_url = 'sqlite:///instance/local.db'
 
-def time_in_beijing():
+def time_in_beijing(strip_time_zone=True):
     '''
     return current time in Beijing as datetime object
     '''
     tz = pytz.timezone('Asia/Shanghai')
     dt = datetime.datetime.now(tz)
+    if strip_time_zone:
+        dt = dt.replace(tzinfo=None)
     return dt
 
+def add_details_to_stock_df(stock_df):
+    '''return df adding sector, aggregate sector, display_name, name to it
+    
+    Parameters
+    ----------
+    stock_df: pd.DataFrame
+        the dataframe contain ticker columns
+    
+    Returns
+    -------
+    merged_df: pd.DataFrame
+        the dataframe with sector, aggregate sector, display_name and name added
+
+    '''
+    with create_engine(db_url).connect() as conn:
+        detail_df = pd.read_sql(ts.STOCKS_DETAILS_TABLE, con=conn)
+        merged_df = pd.merge(stock_df, detail_df[
+            ['sector', 'name',
+             'aggregate_sector',
+             'display_name',
+             'ticker']
+        ], on='ticker', how='left')
+        merged_df['aggregate_sector'].fillna('其他', inplace=True)
+        return merged_df
 
 def convert_string_to_datetime(date_string, time_zone="Asia/Shanghai"):
     '''
@@ -108,3 +137,5 @@ def create_share_changes_report(df):
         markdown += '{} | {} | {}\n'.format(row['ticker'],
                                             row['display_name'], share_changes_str)
     return markdown
+
+