Skip to content

Commit

Permalink
feat: lazy load ansible imports
Browse files Browse the repository at this point in the history
  • Loading branch information
PaulFarault committed Jul 18, 2024
1 parent 0084e7b commit 08d3e55
Show file tree
Hide file tree
Showing 4 changed files with 125 additions and 33 deletions.
4 changes: 2 additions & 2 deletions tdp/cli/commands/default_diff.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,9 @@
from typing import TYPE_CHECKING, Optional

import click
from ansible.utils.vars import merge_hash

from tdp.cli.params import collections_option, vars_option
from tdp.core.ansible_loader import AnsibleLoader
from tdp.core.constants import DEFAULT_VARS_DIRECTORY_NAME
from tdp.core.variables import ClusterVariables, Variables

Expand Down Expand Up @@ -73,7 +73,7 @@ def service_diff(collections, service):
with Variables(default_service_vars_filepath).open(
"r"
) as default_variables:
default_service_varfile = merge_hash(
default_service_varfile = AnsibleLoader.load_merge_hash()(
default_service_varfile, default_variables
)

Expand Down
102 changes: 102 additions & 0 deletions tdp/core/ansible_loader.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
# Copyright 2022 TOSIT.IO
# SPDX-License-Identifier: Apache-2.0


class AnsibleLoader:
"""Lazy loader for Ansible classes and functions.
This class is required as ansible automatically generate a config when imported.
"""

_merge_hash = None
_from_yaml = None
_AnsibleDumper = None
_InventoryCLI = None
_InventoryReader = None
_InventoryManager = None
_CustomInventoryCLI = None

@classmethod
def load_merge_hash(cls):
"""Load the merge_hash function from ansible."""
if cls._merge_hash is None:
from ansible.utils.vars import merge_hash

cls._merge_hash = merge_hash

return cls._merge_hash

@classmethod
def load_from_yaml(cls):
"""Load the from_yaml function from ansible."""
if cls._from_yaml is None:
from ansible.parsing.utils.yaml import from_yaml

cls._from_yaml = from_yaml

return cls._from_yaml

@classmethod
def load_AnsibleDumper(cls):
"""Load the AnsibleDumper class from ansible."""
if cls._AnsibleDumper is None:
from ansible.parsing.yaml.dumper import AnsibleDumper

cls._AnsibleDumper = AnsibleDumper

return cls._AnsibleDumper

@classmethod
def load_InventoryCLI(cls):
"""Load the InventoryCLI class from ansible."""
if cls._InventoryCLI is None:
from ansible.cli.inventory import InventoryCLI

cls._InventoryCLI = InventoryCLI

return cls._InventoryCLI

@classmethod
def load_InventoryReader(cls):
"""Load the InventoryReader class from ansible."""
if cls._InventoryReader is None:
from tdp.core.inventory_reader import InventoryReader

cls._InventoryReader = InventoryReader

return cls._InventoryReader

@classmethod
def load_InventoryManager(cls):
"""Load the InventoryManager class from ansible."""
if cls._InventoryManager is None:
from ansible.inventory.manager import InventoryManager

cls._InventoryManager = InventoryManager

return cls._InventoryManager

@classmethod
def get_CustomInventoryCLI(cls):
if cls._CustomInventoryCLI is None:

class CustomInventoryCLI(cls.load_InventoryCLI()):
"""Represent a custom Ansible inventory CLI which does nothing.
This is used to load inventory files with Ansible code.
"""

def __init__(self):
super().__init__(["program", "--list"])
# "run" must be called from CLI (the parent of InventoryCLI), to
# initialize context (reading ansible.cfg for example).
super(cls.load_InventoryCLI(), self).run()
# Get InventoryManager instance
_, self.inventory, _ = self._play_prereqs()

# Avoid call InventoryCLI "run", we do not need to run InventoryCLI
def run(self):
pass

cls._CustomInventoryCLI = CustomInventoryCLI()

return cls._CustomInventoryCLI
34 changes: 9 additions & 25 deletions tdp/core/inventory_reader.py
Original file line number Diff line number Diff line change
@@ -1,45 +1,29 @@
# Copyright 2022 TOSIT.IO
# SPDX-License-Identifier: Apache-2.0

from typing import Optional, TextIO
from typing import TYPE_CHECKING, Optional, TextIO

import yaml
from ansible.cli.inventory import InventoryCLI
from ansible.inventory.manager import InventoryManager

from tdp.core.ansible_loader import AnsibleLoader

try:
from yaml import CLoader as Loader
except ImportError:
from yaml import Loader


# From ansible/cli/inventory.py
class _CustomInventoryCLI(InventoryCLI):
"""Represent a custom Ansible inventory CLI which does nothing.
This is used to load inventory files with Ansible code.
"""

def __init__(self):
super().__init__(["program", "--list"])
# "run" must be called from CLI (the parent of InventoryCLI), to
# initialize context (reading ansible.cfg for example).
super(InventoryCLI, self).run()
# Get InventoryManager instance
_, self.inventory, _ = self._play_prereqs()

# Avoid call InventoryCLI "run", we do not need to run InventoryCLI
def run(self):
pass


custom_inventory_cli_instance = _CustomInventoryCLI()
if TYPE_CHECKING:
from ansible.inventory.manager import InventoryManager


class InventoryReader:
"""Represent an Ansible inventory reader."""

def __init__(self, inventory: Optional[InventoryManager] = None):
self.inventory = inventory or custom_inventory_cli_instance.inventory
def __init__(self, inventory: Optional["InventoryManager"] = None):
if inventory is None:
inventory = AnsibleLoader.get_CustomInventoryCLI().inventory
self.inventory = inventory

def get_hosts(self, *args, **kwargs) -> list[str]:
"""Takes a pattern or list of patterns and returns a list of matching
Expand Down
18 changes: 12 additions & 6 deletions tdp/core/variables/variables.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,8 @@
from weakref import proxy

import yaml
from ansible.parsing.utils.yaml import from_yaml
from ansible.parsing.yaml.dumper import AnsibleDumper
from ansible.utils.vars import merge_hash

from tdp.core.ansible_loader import AnsibleLoader
from tdp.core.types import PathLike


Expand Down Expand Up @@ -90,7 +88,7 @@ def merge(self, mapping: MutableMapping) -> None:
Args:
mapping: Mapping to merge.
"""
self._content = merge_hash(self._content, mapping)
self._content = AnsibleLoader.load_merge_hash()(self._content, mapping)

def __getitem__(self, key):
return self._content.__getitem__(key)
Expand Down Expand Up @@ -131,7 +129,10 @@ def __init__(self, path: Path, mode: Optional[str] = None):
self._file_path = path
self._file_descriptor = open(self._file_path, mode or "r+")
# Initialize the content of the variables file
super().__init__(content=from_yaml(self._file_descriptor) or {}, name=path.name)
super().__init__(
content=AnsibleLoader.load_from_yaml()(self._file_descriptor) or {},
name=path.name,
)

def __enter__(self) -> _VariablesIOWrapper:
return proxy(self)
Expand All @@ -152,7 +153,12 @@ def _flush_on_disk(self) -> None:
# Write the content of the variables file on disk
self._file_descriptor.seek(0)
self._file_descriptor.write(
yaml.dump(self._content, Dumper=AnsibleDumper, sort_keys=False, width=1000)
yaml.dump(
self._content,
Dumper=AnsibleLoader.load_AnsibleDumper(),
sort_keys=False,
width=1000,
)
)
self._file_descriptor.truncate()
self._file_descriptor.flush()
Expand Down

0 comments on commit 08d3e55

Please sign in to comment.