meerschaum.utils.venv

  1#! /usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3# vim:fenc=utf-8
  4
  5"""
  6Manage virtual environments.
  7"""
  8
  9from __future__ import annotations
 10
 11from meerschaum.utils.typing import Optional, Union, Dict, List, Tuple
 12from meerschaum.utils.threading import RLock, get_ident
 13
 14__all__ = sorted([
 15    'activate_venv', 'deactivate_venv', 'init_venv',
 16    'inside_venv', 'is_venv_active', 'venv_exec',
 17    'venv_executable', 'venv_exists', 'venv_target_path',
 18    'Venv', 'get_venvs', 'verify_venv', 'get_module_venv',
 19])
 20__pdoc__ = {'Venv': True}
 21
 22LOCKS = {
 23    'sys.path': RLock(),
 24    'active_venvs': RLock(),
 25    'venvs_active_counts': RLock(),
 26}
 27
 28active_venvs = set()
 29active_venvs_counts: Dict[str, int] = {}
 30active_venvs_order: List[Optional[str]] = []
 31threads_active_venvs: Dict[int, 'set[str]'] = {}
 32
 33
 34def activate_venv(
 35        venv: Optional[str] = 'mrsm',
 36        color: bool = True,
 37        force: bool = False,
 38        debug: bool = False,
 39        **kw
 40    ) -> bool:
 41    """
 42    Create a virtual environment (if it doesn't exist) and add it to `sys.path` if necessary.
 43
 44    Parameters
 45    ----------
 46    venv: Optional[str], default 'mrsm'
 47        The virtual environment to activate.
 48
 49    color: bool, default True
 50        If `True`, include color in debug text.
 51
 52    force: bool, default False
 53        If `True`, do not exit early even if the venv is currently active.
 54
 55    debug: bool, default False
 56        Verbosity toggle.
 57
 58    Returns
 59    -------
 60    A bool indicating whether the virtual environment was successfully activated.
 61
 62    """
 63    thread_id = get_ident()
 64    if active_venvs_order and active_venvs_order[0] == venv:
 65        if not force:
 66            return True
 67    import sys, platform, os
 68    from meerschaum.config._paths import VIRTENV_RESOURCES_PATH
 69    if debug:
 70        from meerschaum.utils.debug import dprint
 71    if venv is not None:
 72        init_venv(venv=venv, debug=debug)
 73    with LOCKS['active_venvs']:
 74        if thread_id not in threads_active_venvs:
 75            threads_active_venvs[thread_id] = {}
 76        active_venvs.add(venv)
 77        if venv not in threads_active_venvs[thread_id]:
 78            threads_active_venvs[thread_id][venv] = 1
 79        else:
 80            threads_active_venvs[thread_id][venv] += 1
 81
 82        target = venv_target_path(venv, debug=debug).as_posix()
 83        if venv in active_venvs_order:
 84            sys.path.remove(target)
 85            try:
 86                active_venvs_order.remove(venv)
 87            except Exception as e:
 88                pass
 89        if venv is not None:
 90            sys.path.insert(0, target)
 91        else:
 92            if sys.path and sys.path[0] in (os.getcwd(), ''):
 93                sys.path.insert(1, target)
 94            else:
 95                sys.path.insert(0, target)
 96        try:
 97            active_venvs_order.insert(0, venv)
 98        except Exception as e:
 99            pass
100
101    return True
102
103
104def deactivate_venv(
105        venv: str = 'mrsm',
106        color: bool = True,
107        debug: bool = False,
108        previously_active_venvs: Union['set[str]', List[str], None] = None,
109        force: bool = False,
110        **kw
111    ) -> bool:
112    """
113    Remove a virtual environment from `sys.path` (if it's been activated).
114
115    Parameters
116    ----------
117    venv: str, default 'mrsm'
118        The virtual environment to deactivate.
119
120    color: bool, default True
121        If `True`, include color in debug text.
122
123    debug: bool, default False
124        Verbosity toggle.
125
126    previously_active_venvs: Union[Set[str], List[str], None]
127        If provided, skip deactivating if a virtual environment is in this iterable.
128
129    force: bool, default False
130        If `True`, forcibly deactivate the virtual environment.
131        This may cause issues with other threads, so be careful!
132
133    Returns
134    -------
135    Return a bool indicating whether the virtual environment was successfully deactivated.
136
137    """
138    import sys
139    thread_id = get_ident()
140    if venv is None:
141        if venv in active_venvs:
142            active_venvs.remove(venv)
143        return True
144
145    if previously_active_venvs and venv in previously_active_venvs and not force:
146        return True
147
148    with LOCKS['active_venvs']:
149        if venv in threads_active_venvs.get(thread_id, {}):
150            new_count = threads_active_venvs[thread_id][venv] - 1
151            if new_count > 0 and not force:
152                threads_active_venvs[thread_id][venv] = new_count
153                return True
154            else:
155                del threads_active_venvs[thread_id][venv]
156
157        if not force:
158            for other_thread_id, other_venvs in threads_active_venvs.items():
159                if other_thread_id == thread_id:
160                    continue
161                if venv in other_venvs:
162                    return True
163        else:
164            to_delete = [other_thread_id for other_thread_id in threads_active_venvs]
165            for other_thread_id in to_delete:
166                del threads_active_venvs[other_thread_id]
167
168        if venv in active_venvs:
169            active_venvs.remove(venv)
170
171    if sys.path is None:
172        return False
173
174    target = venv_target_path(venv, allow_nonexistent=force, debug=debug).as_posix()
175    with LOCKS['sys.path']:
176        if target in sys.path:
177            sys.path.remove(target)
178            try:
179                active_venvs_order.remove(venv)
180            except Exception as e:
181                pass
182
183    return True
184
185
186def is_venv_active(
187        venv: str = 'mrsm',
188        color : bool = True,
189        debug: bool = False
190    ) -> bool:
191    """
192    Check if a virtual environment is active.
193
194    Parameters
195    ----------
196    venv: str, default 'mrsm'
197        The virtual environment to check.
198
199    color: bool, default True
200        If `True`, include color in debug text.
201
202    debug: bool, default False
203        Verbosity toggle.
204
205    Returns
206    -------
207    A bool indicating whether the virtual environment `venv` is active.
208
209    """
210    return venv in active_venvs
211
212
213verified_venvs = set()
214def verify_venv(
215        venv: str,
216        debug: bool = False,
217    ) -> None:
218    """
219    Verify that the virtual environment matches the expected state.
220    """
221    import pathlib, platform, os, shutil, subprocess, sys
222    from meerschaum.config._paths import VIRTENV_RESOURCES_PATH
223    from meerschaum.utils.process import run_process
224    from meerschaum.utils.misc import make_symlink, is_symlink
225    from meerschaum.utils.warnings import warn
226    venv_path = VIRTENV_RESOURCES_PATH / venv
227    bin_path = venv_path / (
228        'bin' if platform.system() != 'Windows' else "Scripts"
229    )
230    current_python_versioned_name = (
231        'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor)
232        + ('' if platform.system() != 'Windows' else '.exe')
233    )
234
235    if not (bin_path / current_python_versioned_name).exists():
236        init_venv(venv, verify=False, force=True, debug=debug)
237        current_python_in_venv_path = pathlib.Path(venv_executable(venv=venv))
238        current_python_in_sys_path = pathlib.Path(venv_executable(venv=None))
239        if not current_python_in_venv_path.exists():
240            if is_symlink(current_python_in_venv_path):
241                try:
242                    current_python_in_venv_path.unlink()
243                except Exception as e:
244                    print(f"Unable to remove symlink {current_python_in_venv_path}:\n{e}")
245            try:
246                make_symlink(current_python_in_sys_path, current_python_in_venv_path)
247            except Exception as e:
248                print(
249                    f"Unable to create symlink {current_python_in_venv_path} "
250                    + f"to {current_python_in_sys_path}."
251                )
252        files_to_inspect = sorted(os.listdir(bin_path), reverse=True)
253    else:
254        files_to_inspect = [current_python_versioned_name]
255
256    def get_python_version(python_path: pathlib.Path) -> Union[str, None]:
257        """
258        Return the version for the python binary at the given path.
259        """
260        try:
261            ### It might be a broken symlink, so skip on errors.
262            if debug:
263                print(f"Getting python version for {python_path}")
264            proc = run_process(
265                [str(python_path), '-V'],
266                as_proc=True,
267                capture_output=True,
268            )
269            stdout, stderr = proc.communicate(timeout=1.0)
270        except Exception as e:
271            ### E.g. the symlink may be broken.
272            if is_symlink(python_path):
273                try:
274                    python_path.unlink()
275                except Exception as _e:
276                    print(f"Unable to remove broken symlink {python_path}:\n{e}\n{_e}")
277            return None
278        return stdout.decode('utf-8').strip().replace('Python ', '')
279
280    ### Ensure the versions are symlinked correctly.
281    for filename in files_to_inspect:
282        if not filename.startswith('python'):
283            continue
284        python_path = bin_path / filename
285        version = get_python_version(python_path)
286        if version is None:
287            continue
288        try:
289            major_version = version.split('.', maxsplit=1)[0]
290            minor_version = version.split('.', maxsplit=2)[1]
291        except IndexError:
292            return
293        python_versioned_name = (
294            'python' + major_version + '.' + minor_version
295            + ('' if platform.system() != 'Windows' else '.exe')
296        )
297
298        ### E.g. python3.10 actually links to Python 3.10.
299        if filename == python_versioned_name:
300            real_path = pathlib.Path(os.path.realpath(python_path))
301            if not real_path.exists():
302                try:
303                    python_path.unlink()
304                except Exception as e:
305                    pass
306                init_venv(venv, verify=False, force=True, debug=debug)
307                if not python_path.exists():
308                    raise FileNotFoundError(f"Unable to verify Python symlink:\n{python_path}")
309
310            if python_path == real_path:
311                continue
312
313            try:
314                python_path.unlink()
315            except Exception as e:
316                pass
317            success, msg = make_symlink(real_path, python_path)
318            if not success:
319                warn(msg, color=False)
320            continue
321
322        python_versioned_path = bin_path / python_versioned_name
323        if python_versioned_path.exists():
324            ### Avoid circular symlinks.
325            if get_python_version(python_versioned_path) == version:
326                continue
327            python_versioned_path.unlink()
328        shutil.move(python_path, python_versioned_path)
329
330
331tried_virtualenv = False
332def init_venv(
333        venv: str = 'mrsm',
334        verify: bool = True,
335        force: bool = False,
336        debug: bool = False,
337    ) -> bool:
338    """
339    Initialize the virtual environment.
340
341    Parameters
342    ----------
343    venv: str, default 'mrsm'
344        The name of the virtual environment to create.
345
346    verify: bool, default True
347        If `True`, verify that the virtual environment is in the expected state.
348
349    force: bool, default False
350        If `True`, recreate the virtual environment, even if already initalized.
351
352    Returns
353    -------
354    A `bool` indicating success.
355    """
356    if not force and venv in verified_venvs:
357        return True
358    if not force and venv_exists(venv, debug=debug):
359        if verify:
360            verify_venv(venv, debug=debug)
361            verified_venvs.add(venv)
362        return True
363
364    import io
365    from contextlib import redirect_stdout, redirect_stderr
366    import sys, platform, os, pathlib, shutil
367    from meerschaum.config.static import STATIC_CONFIG
368    from meerschaum.config._paths import VIRTENV_RESOURCES_PATH
369    venv_path = VIRTENV_RESOURCES_PATH / venv
370    docker_home_venv_path = pathlib.Path('/home/meerschaum/venvs/mrsm')
371
372    runtime_env_var = STATIC_CONFIG['environment']['runtime']
373    work_dir_env_var = STATIC_CONFIG['environment']['work_dir']
374    if (
375        not force
376        and venv == 'mrsm'
377        and os.environ.get(work_dir_env_var, None) is not None
378        and docker_home_venv_path.exists()
379    ):
380        shutil.move(docker_home_venv_path, venv_path)
381        if verify:
382            verify_venv(venv, debug=debug)
383            verified_venvs.add(venv)
384        return True
385
386    from meerschaum.utils.packages import run_python_package, attempt_import, _get_pip_os_env
387    global tried_virtualenv
388    try:
389        import venv as _venv
390        uv = attempt_import('uv', venv=None, debug=debug)
391        virtualenv = None
392    except ImportError:
393        _venv = None
394        uv = None
395        virtualenv = None
396
397    _venv_success = False
398
399    if uv is not None:
400        _venv_success = run_python_package(
401            'uv',
402            ['venv', venv_path.as_posix(), '-q'],
403            venv = None,
404            env = _get_pip_os_env(),
405            debug = debug,
406        ) == 0
407
408    if _venv is not None and not _venv_success:
409        f = io.StringIO()
410        with redirect_stdout(f):
411            _venv_success = run_python_package(
412                'venv',
413                [venv_path.as_posix()] + (
414                    ['--symlinks'] if platform.system() != 'Windows' else []
415                ),
416                venv=None, debug=debug
417            ) == 0
418        if not _venv_success:
419            print(f"Please install python3-venv.\n{f.getvalue()}\nFalling back to virtualenv...")
420        if not venv_exists(venv, debug=debug):
421            _venv = None
422    if not _venv_success:
423        virtualenv = attempt_import(
424            'virtualenv', venv=None, lazy=False, install=(not tried_virtualenv), warn=False,
425            check_update=False, color=False, debug=debug,
426        )
427        if virtualenv is None:
428            print(
429                "Failed to import `venv` or `virtualenv`! "
430                + "Please install `virtualenv` via pip then restart Meerschaum."
431            )
432            return False
433
434        tried_virtualenv = True
435        try:
436            python_folder = (
437                'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor)
438            )
439            dist_packages_path = (
440                VIRTENV_RESOURCES_PATH /
441                venv / 'local' / 'lib' / python_folder / 'dist-packages'
442            )
443            local_bin_path = VIRTENV_RESOURCES_PATH / venv / 'local' / 'bin'
444            bin_path = VIRTENV_RESOURCES_PATH / venv / 'bin'
445            vtp = venv_target_path(venv=venv, allow_nonexistent=True, debug=debug)
446            if bin_path.exists():
447                try:
448                    shutil.rmtree(bin_path)
449                except Exception as e:
450                    import traceback
451                    traceback.print_exc()
452            virtualenv.cli_run([venv_path.as_posix()])
453            if dist_packages_path.exists():
454                vtp.mkdir(exist_ok=True, parents=True)
455                for file_path in dist_packages_path.glob('*'):
456                    shutil.move(file_path, vtp)
457                shutil.rmtree(dist_packages_path)
458                #  shutil.move(dist_packages_path, vtp)
459                bin_path.mkdir(exist_ok=True, parents=True)
460                for file_path in local_bin_path.glob('*'):
461                    shutil.move(file_path, bin_path)
462                #  shutil.move(local_bin_path, bin_path)
463                shutil.rmtree(local_bin_path)
464
465        except Exception as e:
466            import traceback
467            traceback.print_exc()
468            return False
469    if verify:
470        verify_venv(venv, debug=debug)
471        verified_venvs.add(venv)
472    return True
473
474
475def venv_executable(venv: Optional[str] = 'mrsm') -> str:
476    """
477    The Python interpreter executable for a given virtual environment.
478    """
479    from meerschaum.config._paths import VIRTENV_RESOURCES_PATH
480    import sys, platform, os
481    return (
482        sys.executable if venv is None
483        else str(
484            VIRTENV_RESOURCES_PATH
485            / venv
486            / (
487                'bin' if platform.system() != 'Windows'
488                else 'Scripts'
489            ) / (
490                'python'
491                + str(sys.version_info.major)
492                + '.'
493                + str(sys.version_info.minor)
494            )
495        )
496    )
497
498
499def venv_exec(
500        code: str,
501        venv: Optional[str] = 'mrsm',
502        env: Optional[Dict[str, str]] = None,
503        with_extras: bool = False,
504        as_proc: bool = False,
505        capture_output: bool = True,
506        debug: bool = False,
507    ) -> Union[bool, Tuple[int, bytes, bytes]]:
508    """
509    Execute Python code in a subprocess via a virtual environment's interpeter.
510    Return `True` if the code successfully executes, `False` on failure.
511
512    Parameters
513    ----------
514    code: str
515        The Python code to excecute.
516
517    venv: str, default 'mrsm'
518        The virtual environment to use to get the path for the Python executable.
519        If `venv` is `None`, use the default `sys.executable` path.
520
521    env: Optional[Dict[str, str]], default None
522        Optionally specify the environment variables for the subprocess.
523        Defaults to `os.environ`.
524
525    with_extras: bool, default False
526        If `True`, return a tuple of the exit code, stdout bytes, and stderr bytes.
527
528    as_proc: bool, default False
529        If `True`, return the `subprocess.Popen` object instead of executing.
530
531    Returns
532    -------
533    By default, return a bool indicating success.
534    If `as_proc` is `True`, return a `subprocess.Popen` object.
535    If `with_extras` is `True`, return a tuple of the exit code, stdout bytes, and stderr bytes.
536
537    """
538    import os
539    import subprocess
540    from meerschaum.utils.debug import dprint
541    executable = venv_executable(venv=venv)
542    cmd_list = [executable, '-c', code]
543    if env is None:
544        env = os.environ
545    if debug:
546        dprint(str(cmd_list))
547    if not with_extras and not as_proc:
548        return subprocess.call(cmd_list, env=env) == 0
549
550    stdout, stderr = (None, None) if not capture_output else (subprocess.PIPE, subprocess.PIPE)
551    process = subprocess.Popen(
552        cmd_list,
553        stdout = stdout,
554        stderr = stderr,
555        env = env,
556    )
557    if as_proc:
558        return process
559    stdout, stderr = process.communicate()
560    exit_code = process.returncode
561    return exit_code, stdout, stderr
562
563
564def venv_exists(venv: Union[str, None], debug: bool = False) -> bool:
565    """
566    Determine whether a virtual environment has been created.
567    """
568    target_path = venv_target_path(venv, allow_nonexistent=True, debug=debug)
569    return target_path.exists()
570
571
572def venv_target_path(
573        venv: Union[str, None],
574        allow_nonexistent: bool = False,
575        debug: bool = False,
576    ) -> 'pathlib.Path':
577    """
578    Return a virtual environment's site-package path.
579
580    Parameters
581    ----------
582    venv: Union[str, None]
583        The virtual environment for which a path should be returned.
584
585    allow_nonexistent: bool, default False
586        If `True`, return a path even if it does not exist.
587
588    Returns
589    -------
590    The `pathlib.Path` object for the virtual environment's path.
591
592    """
593    import os, sys, platform, pathlib, site
594    from meerschaum.config._paths import VIRTENV_RESOURCES_PATH
595    from meerschaum.config.static import STATIC_CONFIG
596
597    ### Check sys.path for a user-writable site-packages directory.
598    if venv is None:
599
600        ### Return the known value for the portable environment.
601        environment_runtime = STATIC_CONFIG['environment']['runtime']
602        if os.environ.get(environment_runtime, None) == 'portable':
603            python_version_folder = (
604                'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor)
605            )
606            executable_path = pathlib.Path(sys.executable)
607            site_packages_path = (
608                (
609                    executable_path.parent.parent / 'lib' / python_version_folder / 'site-packages'
610                ) if platform.system() != 'Windows' else (
611                    executable_path.parent / 'Lib' / 'site-packages'
612                )
613            )
614            if not site_packages_path.exists():
615                raise EnvironmentError(f"Could not find '{site_packages_path}'. Does it exist?")
616            return site_packages_path
617
618        if not inside_venv():
619            user_site_packages = site.getusersitepackages()
620            if user_site_packages is None:
621                raise EnvironmentError("Could not determine user site packages.")
622
623            site_path = pathlib.Path(user_site_packages)
624            if not site_path.exists():
625
626                ### Windows does not have `os.geteuid()`.
627                if platform.system() == 'Windows' or os.geteuid() != 0:
628                    site_path.mkdir(parents=True, exist_ok=True)
629                    return site_path
630
631                ### Allow for dist-level paths (running as root).
632                for possible_dist in site.getsitepackages():
633                    dist_path = pathlib.Path(possible_dist)
634                    if not dist_path.exists():
635                        continue
636                    return dist_path
637                
638                raise EnvironmentError("Could not determine the dist-packages directory.")
639
640            return site_path
641
642    venv_root_path = (
643        (VIRTENV_RESOURCES_PATH / venv)
644        if venv is not None else pathlib.Path(sys.prefix)
645    )
646    target_path = venv_root_path
647
648    ### Ensure 'lib' or 'Lib' exists.
649    lib = 'lib' if platform.system() != 'Windows' else 'Lib'
650    if not allow_nonexistent:
651        if not venv_root_path.exists() or lib not in os.listdir(venv_root_path):
652            print(f"Failed to find lib directory for virtual environment '{venv}'.")
653            import traceback
654            traceback.print_stack()
655            sys.exit(1)
656    target_path = target_path / lib
657
658    ### Check if a 'python3.x' folder exists.
659    python_folder = 'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor)
660    if target_path.exists():
661        target_path = (
662            (target_path / python_folder) if python_folder in os.listdir(target_path)
663            else target_path
664        )
665    else:
666        target_path = (
667            (target_path / python_folder) if platform.system() != 'Windows'
668            else target_path
669        )
670
671    ### Ensure 'site-packages' exists.
672    if allow_nonexistent or 'site-packages' in os.listdir(target_path): ### Windows
673        target_path = target_path / 'site-packages'
674    else:
675        import traceback
676        traceback.print_stack()
677        print(f"Failed to find site-packages directory for virtual environment '{venv}'.")
678        print("This may be because you are using a different Python version.")
679        print("Try deleting the following directory and restarting Meerschaum:")
680        print(VIRTENV_RESOURCES_PATH)
681        sys.exit(1)
682
683    return target_path
684
685
686def inside_venv() -> bool:
687    """
688    Determine whether current Python interpreter is running inside a virtual environment.
689    """
690    import sys
691    return (
692        hasattr(sys, 'real_prefix') or (
693            hasattr(sys, 'base_prefix')
694                and sys.base_prefix != sys.prefix
695        )
696    )
697
698
699def get_venvs() -> List[str]:
700    """
701    Return a list of all the virtual environments.
702    """
703    import os
704    from meerschaum.config._paths import VIRTENV_RESOURCES_PATH
705    venvs = []
706    for filename in os.listdir(VIRTENV_RESOURCES_PATH):
707        path = VIRTENV_RESOURCES_PATH / filename
708        if not path.is_dir():
709            continue
710        if not venv_exists(filename):
711            continue
712        venvs.append(filename)
713    return venvs
714
715
716def get_module_venv(module) -> Union[str, None]:
717    """
718    Return the virtual environment where an imported module is installed.
719
720    Parameters
721    ----------
722    module: ModuleType
723        The imported module to inspect.
724
725    Returns
726    -------
727    The name of a venv or `None`.
728    """
729    import pathlib
730    from meerschaum.config.paths import VIRTENV_RESOURCES_PATH
731    module_path = pathlib.Path(module.__file__).resolve()
732    try:
733        rel_path = module_path.relative_to(VIRTENV_RESOURCES_PATH)
734    except ValueError:
735        return None
736
737    return rel_path.as_posix().split('/', maxsplit=1)[0]
738
739
740from meerschaum.utils.venv._Venv import Venv