meerschaum.plugins

   1#! /usr/bin/env python
   2# -*- coding: utf-8 -*-
   3# vim:fenc=utf-8
   4
   5"""
   6Expose plugin management APIs from the `meerschaum.plugins` module.
   7"""
   8
   9from __future__ import annotations
  10
  11import pathlib
  12import functools
  13from collections import defaultdict
  14
  15import meerschaum as mrsm
  16from meerschaum.utils.typing import Callable, Any, Union, Optional, Dict, List, Tuple
  17from meerschaum.utils.threading import RLock
  18from meerschaum.core.Plugin import Plugin
  19
  20_api_plugins: Dict[str, List[Callable[['fastapi.App'], Any]]] = {}
  21_pre_sync_hooks: Dict[Union[str, None], List[Callable[[Any], Any]]] = {}
  22_post_sync_hooks: Dict[Union[str, None], List[Callable[[Any], Any]]] = {}
  23_actions_daemon_enabled: Dict[str, bool] = {}
  24_locks = {
  25    '_api_plugins': RLock(),
  26    '_dash_plugins': RLock(),
  27    '_pre_sync_hooks': RLock(),
  28    '_post_sync_hooks': RLock(),
  29    '_actions_daemon_enabled': RLock(),
  30    '__path__': RLock(),
  31    'sys.path': RLock(),
  32    'internal_plugins': RLock(),
  33    '_synced_symlinks': RLock(),
  34    'PLUGINS_INTERNAL_LOCK_PATH': RLock(),
  35}
  36__all__ = (
  37    "Plugin",
  38    "make_action",
  39    "api_plugin",
  40    "dash_plugin",
  41    "web_page",
  42    "import_plugins",
  43    "from_plugin_import",
  44    "reload_plugins",
  45    "get_plugins",
  46    "get_data_plugins",
  47    "add_plugin_argument",
  48    "pre_sync_hook",
  49    "post_sync_hook",
  50)
  51__pdoc__ = {
  52    'venvs': False,
  53    'data': False,
  54    'stack': False,
  55    'plugins': False,
  56}
  57
  58
  59def make_action(
  60    function: Optional[Callable[[Any], Any]] = None,
  61    shell: bool = False,
  62    activate: bool = True,
  63    deactivate: bool = True,
  64    debug: bool = False,
  65    daemon: bool = True,
  66    skip_if_loaded: bool = True,
  67    _plugin_name: Optional[str] = None,
  68) -> Callable[[Any], Any]:
  69    """
  70    Make a function a Meerschaum action. Useful for plugins that are adding multiple actions.
  71    
  72    Parameters
  73    ----------
  74    function: Callable[[Any], Any]
  75        The function to become a Meerschaum action. Must accept all keyword arguments.
  76        
  77    shell: bool, default False
  78        Not used.
  79        
  80    Returns
  81    -------
  82    Another function (this is a decorator function).
  83
  84    Examples
  85    --------
  86    >>> from meerschaum.plugins import make_action
  87    >>>
  88    >>> @make_action
  89    ... def my_action(**kw):
  90    ...     print('foo')
  91    ...     return True, "Success"
  92    >>>
  93    """
  94    def _decorator(func: Callable[[Any], Any]) -> Callable[[Any], Any]:
  95        from meerschaum.actions import actions, _custom_actions_plugins, _plugins_actions
  96        if skip_if_loaded and func.__name__ in actions:
  97            return func
  98
  99        import meerschaum.config.paths as paths
 100        plugin_name = (
 101            func.__name__.split(
 102                f"{paths.PLUGINS_RESOURCES_PATH.stem}.",
 103                maxsplit=1,
 104            )[-1].split('.')[0]
 105        )
 106        plugin = Plugin(plugin_name) if plugin_name else None
 107
 108        if debug:
 109            from meerschaum.utils.debug import dprint
 110            dprint(
 111                f"Adding action '{func.__name__}' from plugin "
 112                f"'{plugin}'..."
 113            )
 114
 115        actions[func.__name__] = func
 116        _custom_actions_plugins[func.__name__] = plugin_name
 117        if plugin_name not in _plugins_actions:
 118            _plugins_actions[plugin_name] = []
 119        _plugins_actions[plugin_name].append(func.__name__)
 120        if not daemon:
 121            _actions_daemon_enabled[func.__name__] = False
 122        return func
 123
 124    if function is None:
 125        return _decorator
 126    return _decorator(function)
 127
 128
 129def pre_sync_hook(
 130    function: Callable[[Any], Any],
 131) -> Callable[[Any], Any]:
 132    """
 133    Register a function as a sync hook to be executed right before sync.
 134    
 135    Parameters
 136    ----------
 137    function: Callable[[Any], Any]
 138        The function to execute right before a sync.
 139        
 140    Returns
 141    -------
 142    Another function (this is a decorator function).
 143
 144    Examples
 145    --------
 146    >>> from meerschaum.plugins import pre_sync_hook
 147    >>>
 148    >>> @pre_sync_hook
 149    ... def log_sync(pipe, **kwargs):
 150    ...     print(f"About to sync {pipe} with kwargs:\n{kwargs}.")
 151    >>>
 152    """
 153    with _locks['_pre_sync_hooks']:
 154        plugin_name = _get_parent_plugin()
 155        try:
 156            if plugin_name not in _pre_sync_hooks:
 157                _pre_sync_hooks[plugin_name] = []
 158            _pre_sync_hooks[plugin_name].append(function)
 159        except Exception as e:
 160            from meerschaum.utils.warnings import warn
 161            warn(e)
 162    return function
 163
 164
 165def post_sync_hook(
 166    function: Callable[[Any], Any],
 167) -> Callable[[Any], Any]:
 168    """
 169    Register a function as a sync hook to be executed upon completion of a sync.
 170    
 171    Parameters
 172    ----------
 173    function: Callable[[Any], Any]
 174        The function to execute upon completion of a sync.
 175        
 176    Returns
 177    -------
 178    Another function (this is a decorator function).
 179
 180    Examples
 181    --------
 182    >>> from meerschaum.plugins import post_sync_hook
 183    >>> from meerschaum.utils.misc import interval_str
 184    >>> from datetime import timedelta
 185    >>>
 186    >>> @post_sync_hook
 187    ... def log_sync(pipe, success_tuple, duration=None, **kwargs):
 188    ...     duration_delta = timedelta(seconds=duration)
 189    ...     duration_text = interval_str(duration_delta)
 190    ...     print(f"It took {duration_text} to sync {pipe}.")
 191    >>>
 192    """
 193    with _locks['_post_sync_hooks']:
 194        try:
 195            plugin_name = _get_parent_plugin()
 196            if plugin_name not in _post_sync_hooks:
 197                _post_sync_hooks[plugin_name] = []
 198            _post_sync_hooks[plugin_name].append(function)
 199        except Exception as e:
 200            from meerschaum.utils.warnings import warn
 201            warn(e)
 202    return function
 203
 204
 205_plugin_endpoints_to_pages = {}
 206_plugins_web_pages = {}
 207def web_page(
 208    page: Union[str, None, Callable[[Any], Any]] = None,
 209    login_required: bool = True,
 210    skip_navbar: bool = False,
 211    page_group: Optional[str] = None,
 212    **kwargs
 213) -> Any:
 214    """
 215    Quickly add pages to the dash application.
 216
 217    Examples
 218    --------
 219    >>> import meerschaum as mrsm
 220    >>> from meerschaum.plugins import web_page
 221    >>> html = mrsm.attempt_import('dash.html')
 222    >>> 
 223    >>> @web_page('foo/bar', login_required=False)
 224    >>> def foo_bar():
 225    ...     return html.Div([html.H1("Hello, World!")])
 226    >>> 
 227    """
 228    page_str = None
 229
 230    def _decorator(_func: Callable[[Any], Any]) -> Callable[[Any], Any]:
 231        nonlocal page_str, page_group
 232
 233        @functools.wraps(_func)
 234        def wrapper(*_args, **_kwargs):
 235            return _func(*_args, **_kwargs)
 236
 237        if page_str is None:
 238            page_str = _func.__name__
 239
 240        page_str = page_str.lstrip('/').rstrip('/').strip()
 241        if not page_str.startswith('dash'):
 242            page_str = f'/dash/{page_str}'
 243        page_key = (
 244            ' '.join(
 245                [
 246                    word.capitalize()
 247                    for word in (
 248                        page_str.replace('/dash', '').lstrip('/').rstrip('/').strip()
 249                        .replace('-', ' ').replace('_', ' ').split(' ')
 250                    )
 251                ]
 252            )
 253        )
 254 
 255        plugin_name = _get_parent_plugin()
 256        page_group = page_group or plugin_name
 257        if page_group not in _plugin_endpoints_to_pages:
 258            _plugin_endpoints_to_pages[page_group] = {}
 259        _plugin_endpoints_to_pages[page_group][page_str] = {
 260            'function': _func,
 261            'login_required': login_required,
 262            'skip_navbar': skip_navbar,
 263            'page_key': page_key,
 264        }
 265        if plugin_name not in _plugins_web_pages:
 266            _plugins_web_pages[plugin_name] = []
 267        _plugins_web_pages[plugin_name].append(_func)
 268        return wrapper
 269
 270    if callable(page):
 271        decorator_to_return = _decorator(page)
 272        page_str = page.__name__
 273    else:
 274        decorator_to_return = _decorator
 275        page_str = page
 276
 277    return decorator_to_return
 278
 279
 280_dash_plugins = {}
 281def dash_plugin(function: Callable[[Any], Any]) -> Callable[[Any], Any]:
 282    """
 283    Execute the function when starting the Dash application.
 284    """
 285    with _locks['_dash_plugins']:
 286        plugin_name = _get_parent_plugin()
 287        try:
 288            if plugin_name not in _dash_plugins:
 289                _dash_plugins[plugin_name] = []
 290            _dash_plugins[plugin_name].append(function)
 291        except Exception as e:
 292            from meerschaum.utils.warnings import warn
 293            warn(e)
 294    return function
 295
 296
 297def api_plugin(function: Callable[[Any], Any]) -> Callable[[Any], Any]:
 298    """
 299    Execute the function when initializing the Meerschaum API module.
 300    Useful for lazy-loading heavy plugins only when the API is started,
 301    such as when editing the `meerschaum.api.app` FastAPI app.
 302    
 303    The FastAPI app will be passed as the only parameter.
 304    
 305    Examples
 306    --------
 307    >>> from meerschaum.plugins import api_plugin
 308    >>>
 309    >>> @api_plugin
 310    >>> def initialize_plugin(app):
 311    ...     @app.get('/my/new/path')
 312    ...     def new_path():
 313    ...         return {'message': 'It works!'}
 314    >>>
 315    """
 316    with _locks['_api_plugins']:
 317        try:
 318            if function.__module__ not in _api_plugins:
 319                _api_plugins[function.__module__] = []
 320            _api_plugins[function.__module__].append(function)
 321        except Exception as e:
 322            from meerschaum.utils.warnings import warn
 323            warn(e)
 324    return function
 325
 326
 327_synced_symlinks: int = 0
 328_injected_plugin_symlinks = defaultdict(lambda: 0)
 329def sync_plugins_symlinks(debug: bool = False, warn: bool = True) -> None:
 330    """
 331    Update the plugins' internal symlinks. 
 332    """
 333    from meerschaum.utils.warnings import error, warn as _warn, dprint
 334    global _synced_symlinks
 335    with _locks['_synced_symlinks']:
 336        if _synced_symlinks > 1:
 337            if debug:
 338                dprint("Skip syncing symlinks...")
 339            return
 340
 341    import os
 342    import pathlib
 343    import time
 344    from collections import defaultdict
 345    from meerschaum.utils.misc import flatten_list, make_symlink, is_symlink
 346    from meerschaum._internal.static import STATIC_CONFIG
 347    import meerschaum.config.paths as paths
 348
 349    ### If the lock file exists, sleep for up to a second or until it's removed before continuing.
 350    with _locks['PLUGINS_INTERNAL_LOCK_PATH']:
 351        if paths.PLUGINS_INTERNAL_LOCK_PATH.exists():
 352            lock_sleep_total = STATIC_CONFIG['plugins']['lock_sleep_total']
 353            lock_sleep_increment = STATIC_CONFIG['plugins']['lock_sleep_increment']
 354            lock_start = time.perf_counter()
 355            while (
 356                (time.perf_counter() - lock_start) < lock_sleep_total
 357            ):
 358                time.sleep(lock_sleep_increment)
 359                if not paths.PLUGINS_INTERNAL_LOCK_PATH.exists():
 360                    break
 361                try:
 362                    paths.PLUGINS_INTERNAL_LOCK_PATH.unlink()
 363                except Exception as e:
 364                    if warn:
 365                        _warn(
 366                            f"Error while removing lockfile {paths.PLUGINS_INTERNAL_LOCK_PATH}:\n"
 367                            f"{e}"
 368                        )
 369                    break
 370
 371        ### Begin locking from other processes.
 372        try:
 373            paths.PLUGINS_INTERNAL_LOCK_PATH.touch()
 374        except Exception as e:
 375            if warn:
 376                _warn(f"Unable to create lockfile {paths.PLUGINS_INTERNAL_LOCK_PATH}:\n{e}")
 377
 378    with _locks['internal_plugins']:
 379
 380        try:
 381            from importlib.metadata import entry_points
 382        except ImportError:
 383            importlib_metadata = mrsm.attempt_import('importlib_metadata', lazy=False)
 384            entry_points = importlib_metadata.entry_points
 385
 386        ### NOTE: Allow plugins to be installed via `pip`.
 387        packaged_plugin_paths = []
 388        try:
 389            discovered_packaged_plugins_eps = entry_points(group='meerschaum.plugins')
 390        except TypeError:
 391            discovered_packaged_plugins_eps = []
 392
 393        for ep in discovered_packaged_plugins_eps:
 394            module_name = ep.name
 395            for package_file_path in ep.dist.files:
 396                if package_file_path.suffix != '.py':
 397                    continue
 398                if str(package_file_path) == f'{module_name}.py':
 399                    packaged_plugin_paths.append(package_file_path.locate())
 400                elif str(package_file_path) == f'{module_name}/__init__.py':
 401                    packaged_plugin_paths.append(package_file_path.locate().parent)
 402
 403        if is_symlink(paths.PLUGINS_RESOURCES_PATH) or not paths.PLUGINS_RESOURCES_PATH.exists():
 404            try:
 405                paths.PLUGINS_RESOURCES_PATH.unlink()
 406            except Exception:
 407                pass
 408
 409        paths.PLUGINS_RESOURCES_PATH.mkdir(exist_ok=True)
 410
 411        existing_symlinked_paths = {
 412            _existing_symlink: pathlib.Path(os.path.realpath(_existing_symlink))
 413            for item in os.listdir(paths.PLUGINS_RESOURCES_PATH)
 414            if is_symlink(_existing_symlink := (paths.PLUGINS_RESOURCES_PATH / item))
 415        }
 416        injected_symlinked_paths = {
 417            _injected_symlink: pathlib.Path(os.path.realpath(_injected_symlink))
 418            for item in os.listdir(paths.PLUGINS_INJECTED_RESOURCES_PATH)
 419            if is_symlink(_injected_symlink := (paths.PLUGINS_INJECTED_RESOURCES_PATH / item))
 420        }
 421        plugins_to_be_symlinked = list(flatten_list(
 422            [
 423                [
 424                    pathlib.Path(os.path.realpath(plugins_path / item))
 425                    for item in os.listdir(plugins_path)
 426                    if (
 427                        not item.startswith('.')
 428                    ) and (item not in ('__pycache__', '__init__.py'))
 429                ]
 430                for plugins_path in paths.PLUGINS_DIR_PATHS
 431                if plugins_path.exists()
 432            ]
 433        ))
 434        plugins_to_be_symlinked.extend(packaged_plugin_paths)
 435
 436        ### Check for duplicates.
 437        seen_plugins = defaultdict(lambda: 0)
 438        for plugin_path in plugins_to_be_symlinked:
 439            plugin_name = plugin_path.stem
 440            seen_plugins[plugin_name] += 1
 441        for plugin_name, plugin_count in seen_plugins.items():
 442            if plugin_count > 1:
 443                if warn:
 444                    _warn(f"Found duplicate plugins named '{plugin_name}'.")
 445
 446        for plugin_symlink_path, real_path in existing_symlinked_paths.items():
 447
 448            ### Remove invalid symlinks.
 449            if real_path not in plugins_to_be_symlinked:
 450                if _injected_plugin_symlinks[plugin_symlink_path] > 1:
 451                    continue
 452                if plugin_symlink_path in injected_symlinked_paths:
 453                    continue
 454                if real_path in injected_symlinked_paths.values():
 455                    continue
 456                try:
 457                    plugin_symlink_path.unlink()
 458                except Exception:
 459                    pass
 460
 461            ### Remove valid plugins from the to-be-symlinked list.
 462            else:
 463                plugins_to_be_symlinked.remove(real_path)
 464
 465        for plugin_path in plugins_to_be_symlinked:
 466            plugin_symlink_path = paths.PLUGINS_RESOURCES_PATH / plugin_path.name
 467            try:
 468                ### There might be duplicate folders (e.g. __pycache__).
 469                if (
 470                    plugin_symlink_path.exists()
 471                    and
 472                    plugin_symlink_path.is_dir()
 473                    and
 474                    not is_symlink(plugin_symlink_path)
 475                ):
 476                    continue
 477                success, msg = make_symlink(plugin_path, plugin_symlink_path)
 478            except Exception as e:
 479                success, msg = False, str(e)
 480            if not success:
 481                if warn:
 482                    _warn(
 483                        f"Failed to create symlink {plugin_symlink_path} "
 484                        + f"to {plugin_path}:\n    {msg}"
 485                    )
 486
 487    ### Release symlink lock file in case other processes need it.
 488    with _locks['PLUGINS_INTERNAL_LOCK_PATH']:
 489        try:
 490            if paths.PLUGINS_INTERNAL_LOCK_PATH.exists():
 491                paths.PLUGINS_INTERNAL_LOCK_PATH.unlink()
 492        ### Sometimes competing threads will delete the lock file at the same time.
 493        except FileNotFoundError:
 494            pass
 495        except Exception as e:
 496            if warn:
 497                _warn(f"Error cleaning up lockfile {paths.PLUGINS_INTERNAL_LOCK_PATH}:\n{e}")
 498
 499        try:
 500            if not paths.PLUGINS_INIT_PATH.exists():
 501                paths.PLUGINS_INIT_PATH.touch()
 502        except Exception as e:
 503            error(f"Failed to create the file '{paths.PLUGINS_INIT_PATH}':\n{e}")
 504
 505    with _locks['__path__']:
 506        if str(paths.PLUGINS_RESOURCES_PATH.parent) not in __path__:
 507            __path__.append(str(paths.PLUGINS_RESOURCES_PATH.parent))
 508
 509    with _locks['_synced_symlinks']:
 510        _synced_symlinks += 1
 511
 512
 513def import_plugins(
 514    *plugins_to_import: Union[str, List[str], None],
 515    warn: bool = True,
 516) -> Union[
 517    'ModuleType', Tuple['ModuleType', None]
 518]:
 519    """
 520    Import the Meerschaum plugins directory.
 521
 522    Parameters
 523    ----------
 524    plugins_to_import: Union[str, List[str], None]
 525        If provided, only import the specified plugins.
 526        Otherwise import the entire plugins module. May be a string, list, or `None`.
 527        Defaults to `None`.
 528
 529    Returns
 530    -------
 531    A module of list of modules, depening on the number of plugins provided.
 532
 533    """
 534    import sys
 535    import importlib
 536    import meerschaum.config.paths as paths
 537    from meerschaum.utils.misc import flatten_list
 538    from meerschaum.utils.venv import is_venv_active, activate_venv, deactivate_venv, Venv
 539    from meerschaum.utils.warnings import warn as _warn
 540    plugins_to_import = list(plugins_to_import)
 541    prepended_sys_path = False
 542    with _locks['sys.path']:
 543
 544        ### Since plugins may depend on other plugins,
 545        ### we need to activate the virtual environments for library plugins.
 546        ### This logic exists in `Plugin.activate_venv()`,
 547        ### but that code requires the plugin's module to already be imported.
 548        ### It's not a guarantee of correct activation order,
 549        ### e.g. if a library plugin pins a specific package and another 
 550        plugins_names = get_plugins_names()
 551        already_active_venvs = [is_venv_active(plugin_name) for plugin_name in plugins_names]
 552
 553        if not sys.path or sys.path[0] != str(paths.PLUGINS_RESOURCES_PATH.parent):
 554            prepended_sys_path = True
 555            sys.path.insert(0, str(paths.PLUGINS_RESOURCES_PATH.parent))
 556
 557        if not plugins_to_import:
 558            for plugin_name in plugins_names:
 559                activate_venv(plugin_name)
 560            try:
 561                imported_plugins = importlib.import_module(paths.PLUGINS_RESOURCES_PATH.stem)
 562            except ImportError as e:
 563                _warn(f"Failed to import the plugins module:\n    {e}")
 564                import traceback
 565                traceback.print_exc()
 566                imported_plugins = None
 567            for plugin_name in plugins_names:
 568                if plugin_name in already_active_venvs:
 569                    continue
 570                deactivate_venv(plugin_name)
 571
 572        else:
 573            imported_plugins = []
 574            for plugin_name in flatten_list(plugins_to_import):
 575                plugin = Plugin(plugin_name)
 576                try:
 577                    with Venv(plugin, init_if_not_exists=False):
 578                        imported_plugins.append(
 579                            importlib.import_module(
 580                                f'{paths.PLUGINS_RESOURCES_PATH.stem}.{plugin_name}'
 581                            )
 582                        )
 583                except Exception as e:
 584                    _warn(
 585                        f"Failed to import plugin '{plugin_name}':\n    "
 586                        + f"{e}\n\nHere's a stacktrace:",
 587                        stack = False,
 588                    )
 589                    from meerschaum.utils.formatting import get_console
 590                    get_console().print_exception(
 591                        suppress = [
 592                            'meerschaum/plugins/__init__.py',
 593                            importlib,
 594                            importlib._bootstrap,
 595                        ]
 596                    )
 597                    imported_plugins.append(None)
 598
 599        if imported_plugins is None and warn:
 600            _warn("Failed to import plugins.", stacklevel=3)
 601
 602        if prepended_sys_path and str(paths.PLUGINS_RESOURCES_PATH.parent) in sys.path:
 603            sys.path.remove(str(paths.PLUGINS_RESOURCES_PATH.parent))
 604
 605    if isinstance(imported_plugins, list):
 606        return (imported_plugins[0] if len(imported_plugins) == 1 else tuple(imported_plugins))
 607    return imported_plugins
 608
 609
 610def from_plugin_import(plugin_import_name: str, *attrs: str) -> Any:
 611    """
 612    Emulate the `from module import x` behavior.
 613
 614    Parameters
 615    ----------
 616    plugin_import_name: str
 617        The import name of the plugin's module.
 618        Separate submodules with '.' (e.g. 'compose.utils.pipes')
 619
 620    attrs: str
 621        Names of the attributes to return.
 622
 623    Returns
 624    -------
 625    Objects from a plugin's submodule.
 626    If multiple objects are provided, return a tuple.
 627
 628    Examples
 629    --------
 630    >>> init = from_plugin_import('compose.utils', 'init')
 631    >>> with mrsm.Venv('compose'):
 632    ...     cf = init()
 633    >>> build_parent_pipe, get_defined_pipes = from_plugin_import(
 634    ...     'compose.utils.pipes',
 635    ...     'build_parent_pipe',
 636    ...     'get_defined_pipes',
 637    ... )
 638    >>> parent_pipe = build_parent_pipe(cf)
 639    >>> defined_pipes = get_defined_pipes(cf)
 640    """
 641    import importlib
 642    import meerschaum.config.paths as paths
 643    from meerschaum.utils.warnings import warn as _warn
 644    if plugin_import_name.startswith('plugins.'):
 645        plugin_import_name = plugin_import_name[len('plugins.'):]
 646    plugin_import_parts = plugin_import_name.split('.')
 647    plugin_root_name = plugin_import_parts[0]
 648    plugin = mrsm.Plugin(plugin_root_name)
 649
 650    submodule_import_name = '.'.join(
 651        [paths.PLUGINS_RESOURCES_PATH.stem]
 652        + plugin_import_parts
 653    )
 654    if len(attrs) == 0:
 655        raise ValueError(f"Provide which attributes to return from '{submodule_import_name}'.")
 656
 657    attrs_to_return = []
 658    with mrsm.Venv(plugin):
 659        if plugin.module is None:
 660            raise ImportError(f"Unable to import plugin '{plugin}'.")
 661
 662        try:
 663            submodule = importlib.import_module(submodule_import_name)
 664        except ImportError as e:
 665            _warn(
 666                f"Failed to import plugin '{submodule_import_name}':\n    "
 667                + f"{e}\n\nHere's a stacktrace:",
 668                stack=False,
 669            )
 670            from meerschaum.utils.formatting import get_console
 671            get_console().print_exception(
 672                suppress=[
 673                    'meerschaum/plugins/__init__.py',
 674                    importlib,
 675                    importlib._bootstrap,
 676                ]
 677            )
 678            return None
 679
 680        for attr in attrs:
 681            try:
 682                attrs_to_return.append(getattr(submodule, attr))
 683            except Exception:
 684                _warn(f"Failed to access '{attr}' from '{submodule_import_name}'.")
 685                attrs_to_return.append(None)
 686        
 687        if len(attrs) == 1:
 688            return attrs_to_return[0]
 689
 690        return tuple(attrs_to_return)
 691
 692
 693_loaded_plugins: bool = False
 694def load_plugins(
 695    skip_if_loaded: bool = True,
 696    shell: bool = False,
 697    debug: bool = False,
 698) -> None:
 699    """
 700    Import Meerschaum plugins and update the actions dictionary.
 701    """
 702    global _loaded_plugins
 703    from meerschaum.utils.warnings import dprint
 704
 705    if skip_if_loaded and _loaded_plugins:
 706        if debug:
 707            dprint("Skip loading plugins...")
 708        return
 709
 710    from inspect import isfunction, getmembers
 711    import meerschaum.config.paths as paths
 712    from meerschaum.actions import __all__ as _all, modules
 713    from meerschaum.utils.packages import get_modules_from_package
 714
 715    _plugins_names, plugins_modules = get_modules_from_package(
 716        import_plugins(),
 717        names = True,
 718        recursive = True,
 719        modules_venvs = True
 720    )
 721
 722    ### I'm appending here to keep from redefining the modules list.
 723    new_modules = (
 724        [
 725            mod
 726            for mod in modules
 727            if not mod.__name__.startswith(paths.PLUGINS_RESOURCES_PATH.stem + '.')
 728        ]
 729        + plugins_modules
 730    )
 731    n_mods = len(modules)
 732    for mod in new_modules:
 733        modules.append(mod)
 734    for i in range(n_mods):
 735        modules.pop(0)
 736
 737    for module in plugins_modules:
 738        for name, func in getmembers(module):
 739            if not isfunction(func):
 740                continue
 741            if name == module.__name__.split('.')[-1]:
 742                make_action(
 743                    func,
 744                    **{'shell': shell, 'debug': debug},
 745                    _plugin_name=name,
 746                    skip_if_loaded=True,
 747                )
 748
 749    _loaded_plugins = True
 750
 751
 752def unload_custom_actions(plugins: Optional[List[str]] = None, debug: bool = False) -> None:
 753    """
 754    Unload the custom actions added by plugins.
 755    """
 756    from meerschaum.actions import (
 757        actions,
 758        _custom_actions_plugins,
 759        _plugins_actions,
 760    )
 761    from meerschaum._internal.entry import _shell
 762    import meerschaum._internal.shell as shell_pkg
 763
 764    plugins = plugins if plugins is not None else list(_plugins_actions)
 765
 766    for plugin in plugins:
 767        action_names = _plugins_actions.get(plugin, [])
 768        actions_to_remove = {
 769            action_name: actions.get(action_name, None)
 770            for action_name in action_names
 771        }
 772        for action_name in action_names:
 773            _ = actions.pop(action_name, None)
 774            _ = _custom_actions_plugins.pop(action_name, None)
 775            _ = _actions_daemon_enabled.pop(action_name, None)
 776
 777        _ = _plugins_actions.pop(plugin, None)
 778        shell_pkg._remove_shell_actions(
 779            _shell=_shell,
 780            actions=actions_to_remove,
 781        )
 782    
 783
 784def unload_plugins(
 785    plugins: Optional[List[str]] = None,
 786    remove_symlinks: bool = True,
 787    debug: bool = False,
 788) -> None:
 789    """
 790    Unload the specified plugins from memory.
 791    """
 792    global _loaded_plugins, _synced_symlinks
 793    import sys
 794    import meerschaum.config.paths as paths
 795    from meerschaum.connectors import unload_plugin_connectors
 796    if debug:
 797        from meerschaum.utils.warnings import dprint
 798
 799    _loaded_plugins = False
 800    _synced_symlinks = 0
 801
 802    all_plugins = get_plugins_names()
 803    plugins = plugins if plugins is not None else all_plugins
 804    if debug:
 805        dprint(f"Unloading plugins: {plugins}")
 806
 807    unload_custom_actions(plugins, debug=debug)
 808    unload_plugin_connectors(plugins, debug=debug)
 809
 810    module_prefix = f"{paths.PLUGINS_RESOURCES_PATH.stem}."
 811    loaded_modules = [mod_name for mod_name in sys.modules if mod_name.startswith(module_prefix)]
 812
 813    root_plugins_mod = (
 814        sys.modules.get(paths.PLUGINS_RESOURCES_PATH.stem, None)
 815        if sorted(plugins) != sorted(all_plugins)
 816        else sys.modules.pop(paths.PLUGINS_RESOURCES_PATH, None)
 817    )
 818
 819    for plugin_name in plugins:
 820        for mod_name in loaded_modules:
 821            if (
 822                mod_name[len(paths.PLUGINS_RESOURCES_PATH.stem):].startswith(plugin_name + '.')
 823                or mod_name[len(paths.PLUGINS_RESOURCES_PATH.stem):] == plugin_name
 824            ):
 825                _ = sys.modules.pop(mod_name, None)
 826
 827        if root_plugins_mod is not None and plugin_name in root_plugins_mod.__dict__:
 828            try:
 829                delattr(root_plugins_mod, plugin_name)
 830            except Exception:
 831                pass
 832
 833        ### Unload sync hooks.
 834        _ = _pre_sync_hooks.pop(plugin_name, None)
 835        _ = _post_sync_hooks.pop(plugin_name, None)
 836
 837        ### Unload API endpoints and pages.
 838        _ = _dash_plugins.pop(plugin_name, None)
 839        web_page_funcs = _plugins_web_pages.pop(plugin_name, None) or []
 840        page_groups_to_pop = []
 841        for page_group, page_functions in _plugin_endpoints_to_pages.items():
 842            page_functions_to_pop = [
 843                page_str
 844                for page_str, page_payload in page_functions.items()
 845                if page_payload.get('function', None) in web_page_funcs
 846            ]
 847            for page_str in page_functions_to_pop:
 848                page_functions.pop(page_str, None)
 849            if not page_functions:
 850                page_groups_to_pop.append(page_group)
 851        
 852        for page_group in page_groups_to_pop:
 853            _plugin_endpoints_to_pages.pop(page_group, None)
 854
 855        ### Remove all but injected symlinks.
 856        if remove_symlinks:
 857            dir_symlink_path = paths.PLUGINS_RESOURCES_PATH / plugin_name
 858            dir_symlink_injected_path = paths.PLUGINS_INJECTED_RESOURCES_PATH / plugin_name
 859            file_symlink_path = paths.PLUGINS_RESOURCES_PATH / f"{plugin_name}.py"
 860            file_symlink_injected_path = paths.PLUGINS_INJECTED_RESOURCES_PATH / f"{plugin_name}.py"
 861
 862            try:
 863                if dir_symlink_path.exists() and not dir_symlink_injected_path.exists():
 864                    dir_symlink_path.unlink()
 865            except Exception:
 866                pass
 867
 868            try:
 869                if file_symlink_path.exists() and not file_symlink_injected_path.exists():
 870                    file_symlink_path.unlink()
 871            except Exception:
 872                pass
 873
 874
 875def reload_plugins(plugins: Optional[List[str]] = None, debug: bool = False) -> None:
 876    """
 877    Reload plugins back into memory.
 878
 879    Parameters
 880    ----------
 881    plugins: Optional[List[str]], default None
 882        The plugins to reload. `None` will reload all plugins.
 883
 884    """
 885    global _synced_symlinks
 886    unload_plugins(plugins, debug=debug)
 887    _synced_symlinks = 0
 888    sync_plugins_symlinks(debug=debug)
 889    load_plugins(skip_if_loaded=False, debug=debug)
 890
 891
 892def get_plugins(*to_load, try_import: bool = True) -> Union[Tuple[Plugin], Plugin]:
 893    """
 894    Return a list of `Plugin` objects.
 895
 896    Parameters
 897    ----------
 898    to_load:
 899        If specified, only load specific plugins.
 900        Otherwise return all plugins.
 901
 902    try_import: bool, default True
 903        If `True`, allow for plugins to be imported.
 904    """
 905    import meerschaum.config.paths as paths
 906    import os
 907    sync_plugins_symlinks()
 908    _plugins = [
 909        Plugin(name)
 910        for name in (
 911            to_load or [
 912                (
 913                    name if (paths.PLUGINS_RESOURCES_PATH / name).is_dir()
 914                    else name[:-3]
 915                )
 916                for name in os.listdir(paths.PLUGINS_RESOURCES_PATH)
 917                if name != '__init__.py'
 918            ]
 919        )
 920    ]
 921    plugins = tuple(plugin for plugin in _plugins if plugin.is_installed(try_import=try_import))
 922    if len(to_load) == 1:
 923        if len(plugins) == 0:
 924            raise ValueError(f"Plugin '{to_load[0]}' is not installed.")
 925        return plugins[0]
 926    return plugins
 927
 928
 929def get_plugins_names(*to_load, **kw) -> List[str]:
 930    """
 931    Return a list of installed plugins.
 932    """
 933    return [plugin.name for plugin in get_plugins(*to_load, **kw)]
 934
 935
 936def get_plugins_modules(*to_load, **kw) -> List['ModuleType']:
 937    """
 938    Return a list of modules for the installed plugins, or `None` if things break.
 939    """
 940    return [plugin.module for plugin in get_plugins(*to_load, **kw)]
 941
 942
 943def get_data_plugins() -> List[Plugin]:
 944    """
 945    Only return the modules of plugins with either `fetch()` or `sync()` functions.
 946    """
 947    import inspect
 948    plugins = get_plugins()
 949    data_names = {'sync', 'fetch'}
 950    data_plugins = []
 951    for plugin in plugins:
 952        for name, ob in inspect.getmembers(plugin.module):
 953            if not inspect.isfunction(ob):
 954                continue
 955            if name not in data_names:
 956                continue
 957            data_plugins.append(plugin)
 958    return data_plugins
 959
 960
 961def add_plugin_argument(*args, **kwargs) -> None:
 962    """
 963    Add argparse arguments under the 'Plugins options' group.
 964    Takes the same parameters as the regular argparse `add_argument()` function.
 965
 966    Examples
 967    --------
 968    >>> add_plugin_argument('--foo', type=int, help="This is my help text!")
 969    >>> 
 970    """
 971    from meerschaum._internal.arguments._parser import groups, _seen_plugin_args, parser
 972    from meerschaum.utils.warnings import warn
 973    _parent_plugin_name = _get_parent_plugin()
 974    title = f"Plugin '{_parent_plugin_name}' options" if _parent_plugin_name else 'Custom options'
 975    group_key = 'plugin_' + (_parent_plugin_name or '')
 976    if group_key not in groups:
 977        groups[group_key] = parser.add_argument_group(
 978            title = title,
 979        )
 980        _seen_plugin_args[group_key] = set()
 981    try:
 982        if str(args) not in _seen_plugin_args[group_key]:
 983            groups[group_key].add_argument(*args, **kwargs)
 984            _seen_plugin_args[group_key].add(str(args))
 985    except Exception as e:
 986        warn(e)
 987
 988
 989def inject_plugin_path(
 990    plugin_path: pathlib.Path,
 991    plugins_resources_path: Optional[pathlib.Path] = None) -> None:
 992    """
 993    Inject a plugin as a symlink into the internal `plugins` directory.
 994
 995    Parameters
 996    ----------
 997    plugin_path: pathlib.Path
 998        The path to the plugin's source module.
 999    """
1000    import meerschaum.config.paths as paths
1001    from meerschaum.utils.misc import make_symlink
1002    if plugins_resources_path is None:
1003        plugins_resources_path = paths.PLUGINS_RESOURCES_PATH
1004        plugins_injected_resources_path = paths.PLUGINS_INJECTED_RESOURCES_PATH
1005    else:
1006        plugins_injected_resources_path = plugins_resources_path / '.injected'
1007
1008    if plugin_path.is_dir():
1009        plugin_name = plugin_path.name
1010        dest_path = plugins_resources_path / plugin_name
1011        injected_path = plugins_injected_resources_path / plugin_name
1012    elif plugin_path.name == '__init__.py':
1013        plugin_name = plugin_path.parent.name
1014        dest_path = plugins_resources_path / plugin_name
1015        injected_path = plugins_injected_resources_path / plugin_name
1016    elif plugin_path.name.endswith('.py'):
1017        plugin_name = plugin_path.name[:(-1 * len('.py'))]
1018        dest_path = plugins_resources_path / plugin_path.name
1019        injected_path = plugins_injected_resources_path / plugin_path.name
1020    else:
1021        raise ValueError(f"Cannot deduce plugin name from path '{plugin_path}'.")
1022
1023    _injected_plugin_symlinks[dest_path] += 1
1024    make_symlink(plugin_path, dest_path)
1025    make_symlink(plugin_path, injected_path)
1026
1027
1028def _get_parent_plugin(stacklevel: Union[int, Tuple[int, ...]] = (1, 2, 3, 4)) -> Union[str, None]:
1029    """If this function is called from outside a Meerschaum plugin, it will return None."""
1030    import inspect
1031    if not isinstance(stacklevel, tuple):
1032        stacklevel = (stacklevel,)
1033
1034    for _level in stacklevel:
1035        try:
1036            parent_globals = inspect.stack()[_level][0].f_globals
1037            global_name = parent_globals.get('__name__', '')
1038            if global_name.startswith('meerschaum.'):
1039                continue
1040            plugin_name = global_name.replace('plugins.', '').split('.')[0]
1041            if plugin_name.startswith('_') or plugin_name == 'importlib':
1042                continue
1043            return plugin_name
1044        except Exception:
1045            continue
1046
1047    return None