meme / manager.py
donotthink's picture
Upload folder using huggingface_hub
ed4faae verified
Raw History Blame Contribute Delete
3.04 kB
import importlib
import importlib.util
import pkgutil
from datetime import datetime
from pathlib import Path
from typing import Optional, Union
from .config import meme_config
from .exception import NoSuchMeme
from .log import logger
from .meme import CommandShortcut, Meme, MemeArgsType, MemeFunction, MemeParamsType
_memes: dict[str, Meme] = {}
def path_to_module_name(path: Path) -> str:
rel_path = path.resolve().relative_to(Path.cwd().resolve())
if rel_path.stem == "__init__":
return ".".join(rel_path.parts[:-1])
else:
return ".".join(rel_path.parts[:-1] + (rel_path.stem,))
def load_meme(module_path: Union[str, Path]):
module_name = (
path_to_module_name(module_path)
if isinstance(module_path, Path)
else module_path
)
try:
importlib.import_module(module_name)
except Exception as e:
logger.opt(colors=True, exception=e).error(f"Failed to import {module_path}!")
def load_memes(dir_path: Union[str, Path]):
if isinstance(dir_path, Path):
dir_path = str(dir_path.resolve())
for module_info in pkgutil.iter_modules([dir_path]):
if module_info.name.startswith("_"):
continue
if not (
module_spec := module_info.module_finder.find_spec(module_info.name, None)
):
continue
if not (module_path := module_spec.origin):
continue
if not (module_loader := module_spec.loader):
continue
try:
module = importlib.util.module_from_spec(module_spec)
module_loader.exec_module(module)
except Exception as e:
logger.opt(colors=True, exception=e).error(
f"Failed to import {module_path}!"
)
def add_meme(
key: str,
function: MemeFunction,
*,
min_images: int = 0,
max_images: int = 0,
min_texts: int = 0,
max_texts: int = 0,
default_texts: list[str] = [],
args_type: Optional[MemeArgsType] = None,
keywords: list[str] = [],
shortcuts: list[CommandShortcut] = [],
tags: set[str] = set(),
date_created: datetime = datetime(2021, 5, 4),
date_modified: datetime = datetime.now(),
):
if key in _memes:
logger.warning(f'Meme with key "{key}" already exists!')
return
if key in meme_config.meme.meme_disabled_list:
logger.warning(f'The key "{key}" is in the disabled list!')
return
meme = Meme(
key,
function,
MemeParamsType(
min_images, max_images, min_texts, max_texts, default_texts, args_type
),
keywords=keywords,
shortcuts=shortcuts,
tags=tags,
date_created=date_created,
date_modified=date_modified,
)
_memes[key] = meme
def get_meme(key: str) -> Meme:
if key not in _memes:
raise NoSuchMeme(key)
return _memes[key]
def get_memes() -> list[Meme]:
return list(_memes.values())
def get_meme_keys() -> list[str]:
return list(_memes.keys())