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