From 4134881fae838e5a6de7c3f392a3cb775c623a90 Mon Sep 17 00:00:00 2001 From: Jakub Kuczys Date: Tue, 29 Oct 2024 18:58:28 +0100 Subject: [PATCH] Optimize Trivia list loading (#6336) Co-authored-by: Kreusada <67752638+Kreusada@users.noreply.github.com> --- redbot/cogs/trivia/trivia.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/redbot/cogs/trivia/trivia.py b/redbot/cogs/trivia/trivia.py index ea319b06005..b4ab7e5aaed 100644 --- a/redbot/cogs/trivia/trivia.py +++ b/redbot/cogs/trivia/trivia.py @@ -29,6 +29,7 @@ UNIQUE_ID = 0xB3C0E453 _ = Translator("Trivia", __file__) +YAMLSafeLoader = getattr(yaml, "CSafeLoader", yaml.SafeLoader) class InvalidListError(Exception): @@ -759,7 +760,7 @@ async def _save_trivia_list( return buffer = io.BytesIO(await attachment.read()) - trivia_dict = yaml.safe_load(buffer) + trivia_dict = yaml.load(buffer, YAMLSafeLoader) TRIVIA_LIST_SCHEMA.validate(trivia_dict) buffer.seek(0) @@ -803,7 +804,7 @@ def get_core_lists() -> List[pathlib.Path]: return list(core_lists_path.glob("*.yaml")) -def get_list(path: pathlib.Path) -> Dict[str, Any]: +def get_list(path: pathlib.Path, *, validate_schema: bool = True) -> Dict[str, Any]: """ Returns a trivia list dictionary from the given path. @@ -814,12 +815,14 @@ def get_list(path: pathlib.Path) -> Dict[str, Any]: """ with path.open(encoding="utf-8") as file: try: - trivia_dict = yaml.safe_load(file) + trivia_dict = yaml.load(file, YAMLSafeLoader) except yaml.error.YAMLError as exc: raise InvalidListError("YAML parsing failed.") from exc - try: - TRIVIA_LIST_SCHEMA.validate(trivia_dict) - except schema.SchemaError as exc: - raise InvalidListError("The list does not adhere to the schema.") from exc + if validate_schema: + try: + TRIVIA_LIST_SCHEMA.validate(trivia_dict) + except schema.SchemaError as exc: + raise InvalidListError("The list does not adhere to the schema.") from exc + return trivia_dict