From a403d83f7e71efe150c03c740c7a0425f8840389 Mon Sep 17 00:00:00 2001 From: RabbitFoRed <73241991+RabbitFored@users.noreply.github.com> Date: Tue, 10 Sep 2024 14:54:54 +0700 Subject: [PATCH] pyrofork: Ignore excluded plugins currently if you specify a plugin to exclude in pyrogram the client first imports it (https://github.com/Mayuri-Chan/pyrofork/blob/caea59cc17fe96bec40b4c55852ce404c4908fa0/pyrogram/client.py#L874) and add_handler() (https://github.com/Mayuri-Chan/pyrofork/blob/caea59cc17fe96bec40b4c55852ce404c4908fa0/pyrogram/client.py#L880) and then after this it uses remove_handler() (https://github.com/Mayuri-Chan/pyrofork/blob/caea59cc17fe96bec40b4c55852ce404c4908fa0/pyrogram/client.py#L948) this usually works well in most case, but in a few case if the module to exclude has an error, this is not handled and stops the program, so need to fix the modules even if it not the target of interest (as it is in exclude=[]) Co-authored-by: wulan17 Signed-off-by: wulan17 --- pyrogram/client.py | 192 +++++++++++++++++++++++++++++++-------------- 1 file changed, 132 insertions(+), 60 deletions(-) diff --git a/pyrogram/client.py b/pyrogram/client.py index 82c72df6..29873334 100644 --- a/pyrogram/client.py +++ b/pyrogram/client.py @@ -866,39 +866,114 @@ def load_plugins(self): include = plugins.get("include", []) exclude = plugins.get("exclude", []) - count = 0 - - if not include: - for path in sorted(Path(root.replace(".", "/")).rglob("*.py")): - module_path = '.'.join(path.parent.parts + (path.stem,)) - module = import_module(module_path) + exclude_plugins = [] + exclude_handlers = {} - for name in vars(module).keys(): - # noinspection PyBroadException - try: - for handler, group in getattr(module, name).handlers: - if isinstance(handler, Handler) and isinstance(group, int): - self.add_handler(handler, group) + if exclude: + for path, handler in exclude: + module_path = os.path.join( + root.replace(".", "/"), path.replace(".", "/") + ) + if handler is None: + exclude_plugins.append(module_path.replace("/", ".")) + else: + exclude_handlers[module_path.replace("/", ".")] = handler - log.info('[{}] [LOAD] {}("{}") in group {} from "{}"'.format( - self.name, type(handler).__name__, name, group, module_path)) + count = 0 - count += 1 - except Exception: - pass + if not include: + for current_root, dirnames, filenames in os.walk(root.replace(".", "/")): + namespace = current_root.replace("/", ".") + if "__pycache__" in namespace: + continue + if namespace in exclude_plugins: + log.warning( + '[%s] [LOAD] Ignoring namespace "%s"', self.name, namespace + ) + continue + else: + for filename in filenames: + if filename.endswith(".py"): + module_path = namespace + "." + filename[:-3] + if module_path in exclude_plugins: + log.warning( + '[%s] [LOAD] Ignoring namespace "%s"', + self.name, + module_path, + ) + continue + else: + module = import_module(module_path) + + for name in vars(module).keys(): + + # noinspection PyBroadException + try: + for handler, group in getattr( + module, name + ).handlers: + if isinstance( + handler, Handler + ) and isinstance(group, int): + + if ( + module_path in exclude_handlers + and name + in exclude_handlers[module_path] + ): + exclude_handlers[ + module_path + ].remove(name) + log.warning( + '[{}] [LOAD] Ignoring function "{}" from group {} in "{}"'.format( + self.name, + name, + group, + module_path, + ) + ) + continue + + self.add_handler(handler, group) + + log.info( + '[{}] [LOAD] {}("{}") in group {} from "{}"'.format( + self.name, + type(handler).__name__, + name, + group, + module_path, + ) + ) + + count += 1 + except Exception as e: + pass else: for path, handlers in include: - module_path = root.replace("/",".") + "." + path + module_path = root.replace("/", ".") + "." + path + if module_path in exclude_plugins: + log.warning( + '[%s] [LOAD] Ignoring namespace "%s"', self.name, module_path + ) + continue + warn_non_existent_functions = True try: module = import_module(module_path) except ImportError: - log.warning('[%s] [LOAD] Ignoring non-existent module "%s"', self.name, module_path) + log.warning( + '[%s] [LOAD] Ignoring non-existent module "%s"', + self.name, + module_path, + ) continue if "__path__" in dir(module): - log.warning('[%s] [LOAD] Ignoring namespace "%s"', self.name, module_path) + log.warning( + '[%s] [LOAD] Ignoring namespace "%s"', self.name, module_path + ) continue if handlers is None: @@ -910,55 +985,52 @@ def load_plugins(self): try: for handler, group in getattr(module, name).handlers: if isinstance(handler, Handler) and isinstance(group, int): + if ( + module_path in exclude_handlers + and name in exclude_handlers[module_path] + ): + exclude_handlers[module_path].remove(name) + log.warning( + '[{}] [LOAD] Ignoring function "{}" from group {} in "{}"'.format( + self.name, name, group, module_path + ) + ) + continue self.add_handler(handler, group) - log.info('[{}] [LOAD] {}("{}") in group {} from "{}"'.format( - self.name, type(handler).__name__, name, group, module_path)) + log.info( + '[{}] [LOAD] {}("{}") in group {} from "{}"'.format( + self.name, + type(handler).__name__, + name, + group, + module_path, + ) + ) count += 1 except Exception: if warn_non_existent_functions: - log.warning('[{}] [LOAD] Ignoring non-existent function "{}" from "{}"'.format( - self.name, name, module_path)) - - if exclude: - for path, handlers in exclude: - module_path = root.replace("/",".") + "." + path - warn_non_existent_functions = True - - try: - module = import_module(module_path) - except ImportError: - log.warning('[%s] [UNLOAD] Ignoring non-existent module "%s"', self.name, module_path) - continue - - if "__path__" in dir(module): - log.warning('[%s] [UNLOAD] Ignoring namespace "%s"', self.name, module_path) - continue - - if handlers is None: - handlers = vars(module).keys() - warn_non_existent_functions = False - - for name in handlers: - # noinspection PyBroadException - try: - for handler, group in getattr(module, name).handlers: - if isinstance(handler, Handler) and isinstance(group, int): - self.remove_handler(handler, group) - - log.info('[{}] [UNLOAD] {}("{}") from group {} in "{}"'.format( - self.name, type(handler).__name__, name, group, module_path)) + log.warning( + '[{}] [LOAD] Ignoring non-existent function "{}" from "{}"'.format( + self.name, name, module_path + ) + ) - count -= 1 - except Exception: - if warn_non_existent_functions: - log.warning('[{}] [UNLOAD] Ignoring non-existent function "{}" from "{}"'.format( - self.name, name, module_path)) + for module in exclude_handlers: + for handler in exclude_handlers[module]: + log.warning( + '[{}] [LOAD] Ignoring non-existent function "{}" from "{}"'.format( + self.name, handler, module + ) + ) if count > 0: - log.info('[{}] Successfully loaded {} plugin{} from "{}"'.format( - self.name, count, "s" if count > 1 else "", root)) + log.info( + '[{}] Successfully loaded {} plugin{} from "{}"'.format( + self.name, count, "s" if count > 1 else "", root + ) + ) else: log.warning('[%s] No plugin loaded from "%s"', self.name, root)