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