neurostats-API 0.0.21__py3-none-any.whl → 0.0.22__py3-none-any.whl

neurostats_API/fetchers/tech.py

@@ -12,7 +12,9 @@ class TechFetcher(StatsFetcher):
         """
 
         super().__init__(ticker, db_client)
-        self.collection = self.db["TWN/APIPRCD"]
+        if (ticker in self.tw_company_list.keys()):
+            self.collection = self.db["TWN/APIPRCD"]
+
         self.full_ohlcv = self._get_ohlcv()
         self.basic_indexes = [
             'SMA5', 'SMA20', 'SMA60', 'EMA5', 'EMA20', 'EMA40', 'EMA12',
@@ -55,49 +57,31 @@ class TechFetcher(StatsFetcher):
             # 先對yf search
             if self.ticker in ['GSPC', 'IXIC', 'DJI', 'TWII']:
                 full_tick = f'^{self.ticker}'
-            else:
+            elif(self.ticker in self.tw_company_list.keys()):
                 full_tick = f'{self.ticker}.tw'
+            else:
+                full_tick = f"{self.ticker}"
 
             df = self.conduct_yf_search(full_tick)
 
-            if not self.has_required_columns(df, required_cols):
-
-                print(f".tw failed, try .two")
-
+            if (
+                self.ticker in self.tw_company_list.keys() and
+                not self.has_required_columns(df, required_cols)
+            ):
                 full_tick = f'{self.ticker}.two'
 
                 df = self.conduct_yf_search(full_tick)
 
                 if (df.empty):
-                    raise ValueError(f"No data found for ticker: {self.ticker}")
+                    raise ValueError(f"No data found for ticker: {self.ticker}")\
+            
+            return df[required_cols]
 
         except (KeyError, ValueError, TypeError) as e:
-            # 再對TEJ search
-            tej_required_cols = [
-                "mdate", "open_d", 'high_d', 'low_d', 'close_d', 'vol'
-            ]
-            tej_name_proj = {
-                tej_name: org_name
-                for tej_name, org_name in zip(tej_required_cols, required_cols)
-            }
-
-            query = {'ticker': self.ticker}
-            ticker_full = self.collection.find_one(query)
-
-            if not ticker_full:
-                raise ValueError("No ticker found in database")
-
-            daily_data = ticker_full.get("data", [])
-            if not isinstance(daily_data, list):
-                raise TypeError("Expected 'daily_data' to be a list.")
-
-            df = pd.DataFrame(daily_data)
-
-            if not self.has_required_columns(df, tej_required_cols):
-                raise KeyError(f"Missing required columns")
-            df = df.rename(columns=tej_name_proj)
-
-        return df[required_cols]
+            if (self.collection_name in ['TWN/APIPRCD']):
+                return self.conduct_db_search_tej()
+            elif (self.collection_name == 'us_stats'):
+                return self.conduct_db_search_us()
 
     def get_daily(self):
 
@@ -141,6 +125,62 @@ class TechFetcher(StatsFetcher):
         )
 
         return df
+    
+    def conduct_db_search_tej(self):
+        # 再對TEJ search
+        tej_required_cols = [
+            "mdate", "open_d", 'high_d', 'low_d', 'close_d', 'vol'
+        ]
+
+        required_cols = ['date', 'open', 'high', 'low', 'close', 'volume']
+        tej_name_proj = {
+            tej_name: org_name
+            for tej_name, org_name in zip(tej_required_cols, required_cols)
+        }
+
+        query = {'ticker': self.ticker}
+        ticker_full = self.collection.find_one(query)
+
+        if not ticker_full:
+            raise ValueError("No ticker found in database")
+
+        daily_data = ticker_full.get("data", [])
+        if not isinstance(daily_data, list):
+            raise TypeError("Expected 'daily_data' to be a list.")
+
+        df = pd.DataFrame(daily_data)
+
+        if not self.has_required_columns(df, tej_required_cols):
+            raise KeyError(f"Missing required columns")
+        df = df.rename(columns=tej_name_proj)
+
+        return df[required_cols]
+    
+    def conduct_db_search_us(self):
+        required_cols = ['date', 'open', 'high', 'low', 'close', 'volume']
+
+        query = {'ticker': self.ticker}
+        filter_query = {"daily_data" : 1, "_id": 0}
+        ticker_full = self.collection.find_one(query, filter_query)
+
+        if not ticker_full:
+            raise ValueError("No ticker found in database")
+
+        daily_data = ticker_full.get("daily_data", [])
+        if not isinstance(daily_data, list):
+            raise TypeError("Expected 'daily_data' to be a list.")
+
+        df = pd.DataFrame(daily_data)
+
+        if not self.has_required_columns(df, required_cols):
+            missing_cols = [col for col in required_cols if col not in df.columns]
+            missing_cols = ",".join(missing_cols)
+            print(Warning(f"{missing_cols} not in columns"))
+            for col in missing_cols:
+                df[col] = pd.NA
+            # raise KeyError(f"Missing required columns")
+
+        return df[required_cols]
 
 
 class TechProcessor: