Updating profile info works

modules/decorators/callback.py

@@ -1,10 +1,17 @@
+from rich import print
+
 from modules.menu.menu import menu
+from ..exceptions.pretty_exception import PrettyException
 
 
 def callback(f):
     def wrapper(*args, **kwargs):
         menu.pause()
-        f(*args, **kwargs)
+        try:
+            f(*args, **kwargs)
+        except Exception as e:
+            print(PrettyException(e).pretty_exception)
+            input()
         menu.resume()
 
     return wrapper
@@ -14,7 +21,11 @@ def async_callback(f):
     def wrapper(*args, **kwargs):
         import asyncio
         menu.pause()
-        asyncio.run(f(*args, **kwargs))
+        try:
+            asyncio.run(f(*args, **kwargs))
+        except Exception as e:
+            print(PrettyException(e).pretty_exception)
+            input()
         menu.resume()
 
     return wrapper

modules/exceptions/pretty_exception.py

@@ -0,0 +1,47 @@
+import os
+import traceback
+import contextlib
+import re
+
+
+class PrettyException:
+    def __init__(self, e: Exception):
+        self.pretty_exception = f'❌ Error! \n' \
+                                f'🐊 {e.__traceback__.tb_frame.f_code.co_filename.replace(os.getcwd(), "")}' \
+                                f':{e.__traceback__.tb_frame.f_lineno} \n' \
+                                f'😍 {e.__class__.__name__} \n' \
+                                f'👉 {"".join(traceback.format_exception_only(e)).strip()} \n\n' \
+                                f'⬇️ Trace: \n' \
+                                f'{self.get_full_stack()}'
+
+    @staticmethod
+    def get_full_stack():
+        full_stack = traceback.format_exc().replace(
+            "Traceback (most recent call last):\n", ""
+        )
+
+        line_regex = r'  File "(.*?)", line ([0-9]+), in (.+)'
+
+        def format_line(line: str) -> str:
+            filename_, lineno_, name_ = re.search(line_regex, line).groups()
+            with contextlib.suppress(Exception):
+                filename_ = os.path.basename(filename_)
+
+            return (
+                f"🤯 {filename_}:{lineno_} (in"
+                f" {name_} call)"
+            )
+
+        full_stack = "\n".join(
+            [
+                format_line(line)
+                if re.search(line_regex, line)
+                else f"{line}"
+                for line in full_stack.splitlines()
+            ]
+        )
+
+        return full_stack
+
+    def __str__(self):
+        return self.pretty_exception

modules/menu/callbacks/create_session/callbacks.py

@@ -21,6 +21,7 @@ async def create_session_callback():
             username=client.me.username
         ).json()
         print(f'[green]Created[/] session {session_name}...')
+        await client.stop()
         input()
 
     except OperationalError:

modules/menu/callbacks/get_code/__init__.py

@@ -1 +1 @@
-from .get_code import get_code_callback
+from .callbacks import get_code_callback

modules/menu/callbacks/update_info/__init__.py

@@ -0,0 +1 @@
+from .callbacks import update_info_callback

modules/menu/callbacks/update_info/callbacks.py

@@ -0,0 +1,22 @@
+from modules.decorators.callback import async_callback
+from modules.client import GeneratedClient
+from modules.config import sessions
+from modules.config.models import SessionConfig
+from rich import print
+
+
+@async_callback
+async def update_info_callback(session_name: str):
+    print('Getting info about profile...')
+    await (client := GeneratedClient(name=session_name)).start()
+
+    sessions[session_name] = SessionConfig(
+        id=client.me.id,
+        phone=f'+{client.me.phone_number}',
+        profile_name=client.me.first_name + (f' {client.me.last_name}' if client.me.last_name else ''),
+        username=client.me.username
+    ).json()
+    print(f'[green]Updated[/] info for {session_name}. Go back in menu to update it in interface...')
+
+    await client.stop()
+    input()

modules/menu/dynamic/common_generators/valid_sessions_generator.py

@@ -1,21 +1,22 @@
 import os
 from cursesmenu import CursesMenu
 from cursesmenu.items import FunctionItem
+from typing import Callable, Any
 
 from modules.config import sessions
 from modules.config.models import SessionConfig
-from ...callbacks.get_code import get_code_callback
 
 
-def generate_get_code_menu() -> CursesMenu:
+def generate_sessions_list(menu_name: str, callback: Callable[..., Any]) -> CursesMenu:
     submenu = CursesMenu(
-        title='Sessions'
+        title=menu_name,
+        subtitle='name - phone - id - profile name - username'
     )
     for session_name in sessions.keys():
         if os.path.isfile(f'{session_name}.session'):
             submenu.items.append(FunctionItem(
                 f'{session_name} - {SessionConfig(**sessions[session_name])}',
-                function=get_code_callback,
+                function=callback,
                 args=[session_name]
             ))
 

modules/menu/dynamic/get_code/submenu.py

@@ -1,7 +1,8 @@
 from ...menu import menu
 from cursesmenu import CursesMenu
 from modules.custom_items.DynamicSubmenu import DynamicSubmenuItem
-from .generator import generate_get_code_menu
+from ..common_generators.valid_sessions_generator import generate_sessions_list
+from ...callbacks.get_code import get_code_callback
 
 
 get_code_submenu = CursesMenu(
@@ -10,6 +11,7 @@ get_code_submenu = CursesMenu(
 
 get_code_submenu_item = DynamicSubmenuItem(
     text='Get confirmation code',
-    generator=generate_get_code_menu,
+    generator=generate_sessions_list,
+    args=['Get code for', get_code_callback],
     menu=menu
 )

modules/menu/dynamic/update/__init__.py

@@ -0,0 +1 @@
+from .submenu import get_submenu_item

modules/menu/dynamic/update/submenu.py

@@ -0,0 +1,13 @@
+from cursesmenu import CursesMenu
+from modules.custom_items.DynamicSubmenu import DynamicSubmenuItem
+from ..common_generators.valid_sessions_generator import generate_sessions_list
+from ...callbacks.update_info import update_info_callback
+
+
+def get_submenu_item(parent_menu: CursesMenu):
+    return DynamicSubmenuItem(
+        text='Update account info',
+        generator=generate_sessions_list,
+        args=['Update info for', update_info_callback],
+        menu=parent_menu
+    )

modules/menu/items/main.py

@@ -1,10 +1,12 @@
 from .config import config_submenu_item
 from .create_session import create_new_session_item
 from ..dynamic.get_code import get_code_submenu_item
+from .manage import manage_submenu_item
 
 
 items_list = [
     create_new_session_item,
     get_code_submenu_item,
+    manage_submenu_item,
     config_submenu_item
 ]

modules/menu/items/manage/__init__.py

@@ -0,0 +1 @@
+from .submenu import manage_submenu_item

modules/menu/items/manage/fields.py

@@ -0,0 +1,9 @@
+from cursesmenu import CursesMenu
+from ...dynamic import update
+
+
+def get_items_list(menu: CursesMenu):
+    _ = [
+        update.get_submenu_item
+    ]
+    return [x(menu) for x in _]

modules/menu/items/manage/submenu.py

@@ -0,0 +1,17 @@
+from ...menu import menu
+from cursesmenu import CursesMenu
+from cursesmenu.items import SubmenuItem
+from .fields import get_items_list
+
+
+manage_submenu = CursesMenu(
+    title='Manage'
+)
+for item in get_items_list(manage_submenu):
+    manage_submenu.items.append(item)
+
+manage_submenu_item = SubmenuItem(
+    text='Manage/edit/delete sessions',
+    submenu=manage_submenu,
+    menu=menu
+)