diff --git a/src/ansys/dpf/core/custom_operator.py b/src/ansys/dpf/core/custom_operator.py index ea5c53a0af..cdfc5e93cd 100644 --- a/src/ansys/dpf/core/custom_operator.py +++ b/src/ansys/dpf/core/custom_operator.py @@ -8,6 +8,12 @@ import abc import ctypes +import os +import pathlib +import shutil +import warnings +import zipfile + import numpy import traceback @@ -31,6 +37,65 @@ from ansys.dpf.gate import object_handler, capi, dpf_vector, integral_types +def update_virtual_environment_for_custom_operators(restore_original: bool = False): + """Updates the dpf-site.zip file used to start a venv for Python custom operators to run in. + + It updates the site-packages in dpf-site.zip with the site-packages of the current venv. + It stores the original dpf-site.zip for future restoration. + + .. note:: + This feature is only available InProcess to ensure compatibility of the current venv + client-side with the machine where the server is run. + + Parameters + ---------- + restore_original: + If ``True``, restores the original dpf-site.zip. + """ + # Get the path to the dpf-site.zip in the current DPF server + server = dpf.server.get_or_create_server(dpf.SERVER) + if server.has_client(): + raise NotImplementedError( + "Updating the dpf-site.zip of a DPF Server is only available when InProcess." + ) + current_dpf_site_zip_path = os.path.join(server.ansys_path, "dpf", "python", "dpf-site.zip") + # Get the path to where we store the original dpf-site.zip + original_dpf_site_zip_path = os.path.join( + server.ansys_path, "dpf", "python", "original", "dpf-site.zip" + ) + # Restore the original dpf-site.zip + if restore_original: + if os.path.exists(original_dpf_site_zip_path): + shutil.move(src=original_dpf_site_zip_path, dst=current_dpf_site_zip_path) + os.rmdir(os.path.dirname(original_dpf_site_zip_path)) + else: + warnings.warn("No original dpf-site.zip found. Current is most likely the original.") + else: + # Get the current paths to site_packages + import site + paths_to_current_site_packages = site.getsitepackages() + current_site_packages_path = None + # Get the first one targeting an actual site-packages folder + for path_to_site_packages in paths_to_current_site_packages: + if path_to_site_packages[-13:] == "site-packages": + current_site_packages_path = pathlib.Path(path_to_site_packages) + break + if current_site_packages_path is None: + warnings.warn("Could not find a currently loaded site-packages folder to update from.") + return + # Store original dpf-site.zip for this DPF Server if no original is stored + if not os.path.exists(os.path.dirname(original_dpf_site_zip_path)): + os.mkdir(os.path.dirname(original_dpf_site_zip_path)) + shutil.move(src=current_dpf_site_zip_path, dst=original_dpf_site_zip_path) + # Zip the current site-packages at the destination + with zipfile.ZipFile(current_dpf_site_zip_path, mode="w") as archive: + for file_path in current_site_packages_path.rglob("*"): + archive.write( + filename=file_path, + arcname=file_path.relative_to(current_site_packages_path) + ) + + def record_operator(operator_type, *args) -> None: """ Add an operator (with its name, run callback, and specification) to the DPF core registry.