diff --git a/codeflash/languages/java/context.py b/codeflash/languages/java/context.py index 2ccfd34bf..63cc630b0 100644 --- a/codeflash/languages/java/context.py +++ b/codeflash/languages/java/context.py @@ -20,6 +20,8 @@ if TYPE_CHECKING: from tree_sitter import Node +_PARENT_TYPE_NAMES: frozenset[str] = frozenset(("ClassDef", "InterfaceDef", "EnumDef")) + logger = logging.getLogger(__name__) @@ -138,7 +140,7 @@ def _get_parent_type_name(function: FunctionToOptimize) -> str | None: # Check parents for interface/enum if function.parents: for parent in function.parents: - if parent.type in ("ClassDef", "InterfaceDef", "EnumDef"): + if parent.type in _PARENT_TYPE_NAMES: return parent.name return None