Add delete tickers functionality

Author: DUYN

bot/bot.py

@@ -2,6 +2,7 @@
 
 from typing import Any, Dict
 
+from bot.data import remove_ticker, get_tickers, add_ticker, get_company_profile
 from bot import __version__, OperationalException
 from bot.data.data_provider_manager import DataProviderManager
 from bot.data.data_providers import DataProviderException
@@ -14,9 +15,7 @@ class Bot:
     def __init__(self, config: Dict[str, Any]) -> None:
         logger.info('Starting bot version %s', __version__)
 
-        # Init objects
-
-        # Make variable private, only bot should change them
+        # Make variables private, only bot should change them
         self.__config = config
         self.__data_provider_manager = DataProviderManager(self.config)
 
@@ -28,11 +27,54 @@ def set_config(self, config: Dict[str, Any]):
         self.__config = config
 
     def add_ticker(self, ticker: str) -> None:
+        logger.info("Adding ticker ...")
+
+        if not get_company_profile(ticker, self.config):
+            profile = self.__data_provider_manager.get_profile(ticker)
+
+            if not profile:
+                raise OperationalException("Could not evaluate {} with the data providers".format(ticker))
+
+            company_name = profile.get('profile', {}).get('companyName', None)
+            category = profile.get('profile', {}).get('industry', None)
+
+            if not company_name:
+                raise OperationalException("Could not evaluate company name for ticker {} with the data providers")
+
+            if not company_name:
+                raise OperationalException("Could not evaluate category for ticker {} with the data providers")
+
+            try:
+                add_ticker(
+                    ticker,
+                    company_name=company_name,
+                    category=category,
+                    config=self.config
+                )
+            except Exception:
+                raise OperationalException(
+                    "Something went wrong with adding ticker {} to the registry".format(ticker)
+                )
+        else:
+            raise OperationalException(
+                "Ticker {} is already present in registry".format(ticker)
+            )
+
+    def remove_ticker(self, ticker: str) -> None:
+        logger.info("Removing ticker ...")
+
+        if get_company_profile(ticker, self.config):
+
+            try:
+                remove_ticker(ticker, self.config)
+            except Exception:
+                raise OperationalException("Something went wrong while deleting the ticker from the registry")
+        else:
+            raise OperationalException("Provided ticker {} does not exist".format(ticker))
 
-        try:
-            self.__data_provider_manager.add_ticker(ticker)
-        except DataProviderException as e:
-            raise OperationalException(str(e))
+    def list_tickers(self):
+        logger.info("Listing tickers ...")
+        return get_tickers(self.config)
 
 
 

bot/data/__init__.py

@@ -1,6 +1,6 @@
 import logging
 import sqlite3
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Tuple
 
 from bot import setup
 
@@ -31,7 +31,7 @@ def create_tables(config: Dict[str, Any]) -> None:
         con.close()
 
     if TRADES_TABLE_NAME not in table_names:
-        logger.info("Creating database table {}...".format(TRADES_TABLE_NAME))
+        logger.info("Creating database table {} ...".format(TRADES_TABLE_NAME))
         con = create_connection(config)
 
         # Create open trades table
@@ -45,44 +45,65 @@ def create_tables(config: Dict[str, Any]) -> None:
 
 
 def add_ticker(ticker: str, company_name: str, category: str, config: Dict[str, Any]) -> None:
+
+    logger.info("Adding ticker {} to registry...".format(ticker))
+
     con = create_connection(config)
     cursor = con.cursor()
 
-    logger.info("Adding ticker {} ...".format(ticker))
+    # Add ticker if not exists
+    insert_statement = """
+           INSERT INTO TICKERS (ticker, company_name, category)
+           VALUES (?, ?, ?);
+       """
 
-    # Get ticker if exists
-    select_statement = """SELECT ticker_id from TICKERS where ticker = ?"""
-    cursor.execute(select_statement, (ticker,))
+    data_tuple = (ticker, company_name, category)
+    cursor.execute(insert_statement, data_tuple)
 
-    result = cursor.fetchall()
+    con.commit()
+    con.close()
 
-    if result:
-        logger.info("Ticker already in database")
-    else:
-        # Add ticker if not exists
-        insert_statement = """
-               INSERT INTO TICKERS (ticker, company_name, category)
-               VALUES (?, ?, ?);
-           """
 
-        data_tuple = (ticker, company_name, category)
-        cursor.execute(insert_statement, data_tuple)
+def remove_ticker(ticker: str, config: Dict[str, Any]) -> None:
+
+    logger.info("Removing ticker {} from registry ...".format(ticker))
+    con = create_connection(config)
+    cursor = con.cursor()
+
+    delete_statement = '''
+        DELETE from TICKERS where ticker = ? 
+    '''
 
+    cursor.execute(delete_statement, (ticker,))
     con.commit()
     con.close()
 
-def remove_ticker(ticker: str, config: Dict[str, Any]) -> None:
-    pass
 
-def get_company_info(ticker: str, config: Dict[str, any]) -> List[str]:
+def get_company_profile(ticker: str, config: Dict[str, any]) -> Tuple[str]:
     con = create_connection(config)
     cursor = con.cursor()
 
-    logger.info("Getting {} company info ...".format(ticker))
+    logger.info("Getting {} company info from registry ...".format(ticker))
+
+    cursor.execute(" SELECT * FROM TICKERS WHERE ticker=?", (ticker, ))
+    result = cursor.fetchall()
+
+    con.close()
+    return result
+
+
+def get_tickers(config: Dict[str, any]) -> List[str]:
+    con = create_connection(config)
+    cursor = con.cursor()
+
+    logger.info("Getting all tickers from registry ...")
 
     # Get ticker if exists
-    select_statement = """SELECT ticker_id from TICKERS where ticker = ?"""
-    cursor.execute(select_statement, (ticker,))
+    select_statement = '''
+       SELECT (ticker) from TICKERS
+    '''
+
+    cursor.execute(select_statement)
     result = cursor.fetchall()
     con.close()
     return result

bot/data/data_provider_manager.py

@@ -2,7 +2,7 @@
 from typing import List, Dict, Any
 
 from bot.data.data_providers import DataProvider, DataProviderException
-from . import add_ticker, get_all_table_names, create_tables, get_company_info
+from . import get_all_table_names, create_tables
 
 logger = logging.getLogger(__name__)
 
@@ -29,26 +29,23 @@ def __init__(self, config: Dict[str, Any]) -> None:
             from bot.data.data_providers.fmp_data_provider import FMPDataProvider
             self.registered_modules.append(FMPDataProvider(self.__config))
 
-    def add_ticker(self, ticker: str) -> None:
+    def evaluate_ticker(self, ticker: str) -> bool:
 
-        if get_company_info(ticker, self.__config):
-            raise DataProviderException("Ticker {} already exists".format(ticker))
-        else:
-            for data_provider in self.registered_modules:
+        for data_provider in self.registered_modules:
 
-                if data_provider.evaluate_ticker(ticker):
-                    logger.info("Ticker exists")
-                    profile = data_provider.get_profile(ticker)
+            if data_provider.evaluate_ticker(ticker):
+                logger.info("Ticker exists")
+                return True
 
-                    company_name = profile.get('profile', {}).get('companyName', {})
-                    industry = profile.get('profile', {}).get('industry', {})
+        return False
 
-                    if company_name and industry:
-                        add_ticker(ticker, company_name, industry, self.__config)
-                        logger.info("Ticker {} has been added ...".format(ticker))
-                        return
+    def get_profile(self, ticker: str) -> Dict:
 
-        raise DataProviderException("Could not evaluate ticker {}".format(ticker))
+        for data_provider in self.registered_modules:
 
+            profile = data_provider.get_profile(ticker)
 
+            if profile:
+                return profile
 
+        raise DataProviderException("Could not profile for {}".format(ticker))

bot/data/data_providers/fmp_data_provider.py

@@ -11,6 +11,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 class FMPDataProvider(DataProvider):
 
     def __init__(self, config: Dict[str, Any]):

bot/services/service.py

@@ -48,3 +48,6 @@ def cleanup(self) -> None:
 
     def _add_ticker(self, ticker: str) -> None:
         self._bot.add_ticker(ticker)
+
+    def _remove_ticker(self, ticker: str) -> None:
+        self._bot.remove_ticker(ticker)

bot/services/telegram.py

@@ -16,14 +16,21 @@
 
 # Telegram keyboard buttons
 DEFAULT_KEYBOARD_BUTTONS = [
-    ['/add_tickers'],
+    ['/list_tickers', '/add_or_remove_tickers'],
     ['/help', '/version']
 ]
 
-ADDING_TICKERS_CONVERSATION_BUTTONS = [
+STANDARD_CONVERSATION_BUTTONS = [
     ['/cancel']
 ]
 
+TICKERS_CONVERSATION_BUTTONS = [
+    ['/add_tickers', '/remove_tickers', '/cancel']
+]
+
+# Conversation states
+ADDING, REMOVING, LISTING_TICKERS, CHOOSING = range(4)
+
 
 def authorized_only(command_handler: Callable[..., None]) -> Callable[..., Any]:
     """
@@ -61,10 +68,6 @@ def wrapper(self, *args, **kwargs):
 class Telegram(Service):
     """  This class handles all telegram communication """
 
-    # Conversation states
-    class AddTickersConversationState:
-        ADDING, TYPING_REPLY, TYPING_CHOICE = range(3)
-
     def __init__(self, bot) -> None:
         """
         Init the Telegram call, and init the super class service
@@ -82,9 +85,14 @@ def startup(self) -> None:
 
         # States of adding a ticker
         ticker_conversation_handler = ConversationHandler(
-            entry_points=[CommandHandler('add_tickers', self._start_adding_tickers)],
+            entry_points=[CommandHandler('add_or_remove_tickers', self._start_tickers_conversation)],
             states={
-                self.AddTickersConversationState.ADDING: [MessageHandler(Filters.text, self._add_tickers)],
+                CHOOSING: [
+                    CommandHandler('add_tickers', self._start_adding_tickers),
+                    CommandHandler('remove_tickers', self._start_removing_tickers),
+                ],
+                ADDING: [MessageHandler(Filters.text, self._add_tickers)],
+                REMOVING: [MessageHandler(Filters.text, self._remove_tickers)],
             },
             fallbacks=[CommandHandler('cancel', self._cancel_conversation)]
         )
@@ -93,6 +101,7 @@ def startup(self) -> None:
         handles = [
             CommandHandler('help', self._help),
             CommandHandler('version', self._version),
+            CommandHandler('list_tickers', self._list_tickers),
             ticker_conversation_handler
         ]
 
@@ -136,7 +145,8 @@ def _help(self, update: Update, context: CallbackContext) -> None:
 
         message = "*/help:* `This help message`\n" \
                   "*/version:* `Show version`\n" \
-                  "*/add_tickers:* `Add tickers to the registry, so they can be analyzed.`\n"
+                  "*/add_or_remove_tickers:* `Add or remove tickers to the registry`\n" \
+                  "*/list_tickers:* `List all saved tickers in the registry`\n"
 
         self._send_msg(message)
 
@@ -192,11 +202,22 @@ def _version(self, update: Update, context: CallbackContext) -> None:
         """
         self._send_msg('*Version:* `{}`'.format(__version__))
 
+    @authorized_only
+    def _start_tickers_conversation(self, update: Update, context: CallbackContext):
+        self._send_msg("Make your choice", keyboard_buttons=TICKERS_CONVERSATION_BUTTONS)
+        return CHOOSING
+
     @authorized_only
     def _start_adding_tickers(self, update: Update, context: CallbackContext):
         self._send_msg("Please provide the tickers separated by commas, if you submit "
-                       "one ticker you can leave out the comma", keyboard_buttons=ADDING_TICKERS_CONVERSATION_BUTTONS)
-        return self.AddTickersConversationState.ADDING
+                       "one ticker you can leave out the comma", keyboard_buttons=STANDARD_CONVERSATION_BUTTONS)
+        return ADDING
+
+    @authorized_only
+    def _start_removing_tickers(self, update: Update, context: CallbackContext):
+        self._send_msg("Please provide the tickers separated by commas, if you submit "
+                       "one ticker you can leave out the comma", keyboard_buttons=STANDARD_CONVERSATION_BUTTONS)
+        return REMOVING
 
     @authorized_only
     def _add_tickers(self, update: Update, context: CallbackContext):
@@ -221,9 +242,30 @@ def _add_tickers(self, update: Update, context: CallbackContext):
     def _remove_tickers(self, update: Update, context: CallbackContext):
         text = update.message.text
         tickers = [ticker.strip() for ticker in text.split(',')]
-        added_tickers = []
+        removed_tickers = []
+
+        for ticker in tickers:
+
+            try:
+                self._remove_ticker(ticker)
+                removed_tickers.append(ticker)
+            except OperationalException as e:
+                self._send_msg(str(e))
+
+        if removed_tickers:
+            self._send_msg("{} removed".format(removed_tickers))
+
         return ConversationHandler.END
 
+    @authorized_only
+    def _list_tickers(self, update: Update, context: CallbackContext):
+
+        try:
+            tickers = self._bot.list_tickers()
+            self._send_msg(tickers)
+        except Exception as e:
+            self._send_msg(str(e))
+
     @authorized_only
     def _cancel_conversation(self, update: Update, context: CallbackContext):
         logger.info("Conversation is canceled")