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',
 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 = str(venv_target_path(venv, debug=debug))
 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 = str(venv_target_path(venv, allow_nonexistent=force, debug=debug))
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 sys, platform, os, pathlib, shutil
365    from meerschaum.config.static import STATIC_CONFIG
366    from meerschaum.config._paths import VIRTENV_RESOURCES_PATH
367    venv_path = VIRTENV_RESOURCES_PATH / venv
368    docker_home_venv_path = pathlib.Path('/home/meerschaum/venvs/mrsm')
369
370    runtime_env_var = STATIC_CONFIG['environment']['runtime']
371    work_dir_env_var = STATIC_CONFIG['environment']['work_dir']
372    if (
373        not force
374        and venv == 'mrsm'
375        and os.environ.get(work_dir_env_var, None) is not None
376        and docker_home_venv_path.exists()
377    ):
378        shutil.move(docker_home_venv_path, venv_path)
379        if verify:
380            verify_venv(venv, debug=debug)
381            verified_venvs.add(venv)
382        return True
383
384    from meerschaum.utils.packages import run_python_package, attempt_import
385    global tried_virtualenv
386    try:
387        import venv as _venv
388        virtualenv = None
389    except ImportError:
390        _venv = None
391        virtualenv = None
392    
393
394    _venv_success = False
395    if _venv is not None:
396        import io
397        from contextlib import redirect_stdout
398        f = io.StringIO()
399        with redirect_stdout(f):
400            _venv_success = run_python_package(
401                'venv',
402                [str(venv_path)] + (
403                    ['--symlinks'] if platform.system() != 'Windows' else []
404                ),
405                venv=None, debug=debug
406            ) == 0
407        if not _venv_success:
408            print(f"Please install python3-venv.\n{f.getvalue()}\nFalling back to virtualenv...")
409        if not venv_exists(venv, debug=debug):
410            _venv = None
411    if not _venv_success:
412        virtualenv = attempt_import(
413            'virtualenv', venv=None, lazy=False, install=(not tried_virtualenv), warn=False,
414            check_update=False, color=False, debug=debug,
415        )
416        if virtualenv is None:
417            print(
418                "Failed to import `venv` or `virtualenv`! "
419                + "Please install `virtualenv` via pip then restart Meerschaum."
420            )
421            return False
422
423        tried_virtualenv = True
424        try:
425            python_folder = (
426                'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor)
427            )
428            dist_packages_path = (
429                VIRTENV_RESOURCES_PATH /
430                venv / 'local' / 'lib' / python_folder / 'dist-packages'
431            )
432            local_bin_path = VIRTENV_RESOURCES_PATH / venv / 'local' / 'bin'
433            bin_path = VIRTENV_RESOURCES_PATH / venv / 'bin'
434            vtp = venv_target_path(venv=venv, allow_nonexistent=True, debug=debug)
435            if bin_path.exists():
436                try:
437                    shutil.rmtree(bin_path)
438                except Exception as e:
439                    import traceback
440                    traceback.print_exc()
441            virtualenv.cli_run([str(venv_path)])
442            if dist_packages_path.exists():
443                vtp.mkdir(exist_ok=True, parents=True)
444                for file_path in dist_packages_path.glob('*'):
445                    shutil.move(file_path, vtp)
446                shutil.rmtree(dist_packages_path)
447                #  shutil.move(dist_packages_path, vtp)
448                bin_path.mkdir(exist_ok=True, parents=True)
449                for file_path in local_bin_path.glob('*'):
450                    shutil.move(file_path, bin_path)
451                #  shutil.move(local_bin_path, bin_path)
452                shutil.rmtree(local_bin_path)
453
454        except Exception as e:
455            import traceback
456            traceback.print_exc()
457            return False
458    if verify:
459        verify_venv(venv, debug=debug)
460        verified_venvs.add(venv)
461    return True
462
463
464def venv_executable(venv: Optional[str] = 'mrsm') -> str:
465    """
466    The Python interpreter executable for a given virtual environment.
467    """
468    from meerschaum.config._paths import VIRTENV_RESOURCES_PATH
469    import sys, platform, os
470    return (
471        sys.executable if venv is None
472        else str(
473            VIRTENV_RESOURCES_PATH
474            / venv
475            / (
476                'bin' if platform.system() != 'Windows'
477                else 'Scripts'
478            ) / (
479                'python'
480                + str(sys.version_info.major)
481                + '.'
482                + str(sys.version_info.minor)
483            )
484        )
485    )
486
487
488def venv_exec(
489        code: str,
490        venv: Optional[str] = 'mrsm',
491        env: Optional[Dict[str, str]] = None,
492        with_extras: bool = False,
493        as_proc: bool = False,
494        capture_output: bool = True,
495        debug: bool = False,
496    ) -> Union[bool, Tuple[int, bytes, bytes]]:
497    """
498    Execute Python code in a subprocess via a virtual environment's interpeter.
499    Return `True` if the code successfully executes, `False` on failure.
500
501    Parameters
502    ----------
503    code: str
504        The Python code to excecute.
505
506    venv: str, default 'mrsm'
507        The virtual environment to use to get the path for the Python executable.
508        If `venv` is `None`, use the default `sys.executable` path.
509
510    env: Optional[Dict[str, str]], default None
511        Optionally specify the environment variables for the subprocess.
512        Defaults to `os.environ`.
513
514    with_extras: bool, default False
515        If `True`, return a tuple of the exit code, stdout bytes, and stderr bytes.
516
517    as_proc: bool, default False
518        If `True`, return the `subprocess.Popen` object instead of executing.
519
520    Returns
521    -------
522    By default, return a bool indicating success.
523    If `as_proc` is `True`, return a `subprocess.Popen` object.
524    If `with_extras` is `True`, return a tuple of the exit code, stdout bytes, and stderr bytes.
525
526    """
527    import os
528    import subprocess
529    from meerschaum.utils.debug import dprint
530    executable = venv_executable(venv=venv)
531    cmd_list = [executable, '-c', code]
532    if env is None:
533        env = os.environ
534    if debug:
535        dprint(str(cmd_list))
536    if not with_extras and not as_proc:
537        return subprocess.call(cmd_list, env=env) == 0
538
539    stdout, stderr = (None, None) if not capture_output else (subprocess.PIPE, subprocess.PIPE)
540    process = subprocess.Popen(
541        cmd_list,
542        stdout = stdout,
543        stderr = stderr,
544        env = env,
545    )
546    if as_proc:
547        return process
548    stdout, stderr = process.communicate()
549    exit_code = process.returncode
550    return exit_code, stdout, stderr
551
552
553def venv_exists(venv: Union[str, None], debug: bool = False) -> bool:
554    """
555    Determine whether a virtual environment has been created.
556    """
557    target_path = venv_target_path(venv, allow_nonexistent=True, debug=debug)
558    return target_path.exists()
559
560
561def venv_target_path(
562        venv: Union[str, None],
563        allow_nonexistent: bool = False,
564        debug: bool = False,
565    ) -> 'pathlib.Path':
566    """
567    Return a virtual environment's site-package path.
568
569    Parameters
570    ----------
571    venv: Union[str, None]
572        The virtual environment for which a path should be returned.
573
574    allow_nonexistent: bool, default False
575        If `True`, return a path even if it does not exist.
576
577    Returns
578    -------
579    The `pathlib.Path` object for the virtual environment's path.
580
581    """
582    import os, sys, platform, pathlib, site
583    from meerschaum.config._paths import VIRTENV_RESOURCES_PATH
584    from meerschaum.config.static import STATIC_CONFIG
585
586    ### Check sys.path for a user-writable site-packages directory.
587    if venv is None:
588
589        ### Return the known value for the portable environment.
590        environment_runtime = STATIC_CONFIG['environment']['runtime']
591        if os.environ.get(environment_runtime, None) == 'portable':
592            python_version_folder = (
593                'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor)
594            )
595            executable_path = pathlib.Path(sys.executable)
596            site_packages_path = (
597                (
598                    executable_path.parent.parent / 'lib' / python_version_folder / 'site-packages'
599                ) if platform.system() != 'Windows' else (
600                    executable_path.parent / 'Lib' / 'site-packages'
601                )
602            )
603            if not site_packages_path.exists():
604                raise EnvironmentError(f"Could not find '{site_packages_path}'. Does it exist?")
605            return site_packages_path
606
607        if not inside_venv():
608            site_path = pathlib.Path(site.getusersitepackages())
609            if not site_path.exists():
610
611                ### Windows does not have `os.geteuid()`.
612                if platform.system() == 'Windows' or os.geteuid() != 0:
613                    site_path.mkdir(parents=True, exist_ok=True)
614                    return site_path
615
616                ### Allow for dist-level paths (running as root).
617                for possible_dist in reversed(site.getsitepackages()):
618                    dist_path = pathlib.Path(possible_dist)
619                    if not dist_path.exists():
620                        continue
621                    return dist_path
622                
623                raise EnvironmentError("Could not determine the dist-packages directory.")
624
625            return site_path
626
627    venv_root_path = (
628        (VIRTENV_RESOURCES_PATH / venv)
629        if venv is not None else pathlib.Path(sys.prefix)
630    )
631    target_path = venv_root_path
632
633    ### Ensure 'lib' or 'Lib' exists.
634    lib = 'lib' if platform.system() != 'Windows' else 'Lib'
635    if not allow_nonexistent:
636        if not venv_root_path.exists() or lib not in os.listdir(venv_root_path):
637            print(f"Failed to find lib directory for virtual environment '{venv}'.")
638            import traceback
639            traceback.print_stack()
640            sys.exit(1)
641    target_path = target_path / lib
642
643    ### Check if a 'python3.x' folder exists.
644    python_folder = 'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor)
645    if target_path.exists():
646        target_path = (
647            (target_path / python_folder) if python_folder in os.listdir(target_path)
648            else target_path
649        )
650    else:
651        target_path = (
652            (target_path / python_folder) if platform.system() != 'Windows'
653            else target_path
654        )
655
656    ### Ensure 'site-packages' exists.
657    if allow_nonexistent or 'site-packages' in os.listdir(target_path): ### Windows
658        target_path = target_path / 'site-packages'
659    else:
660        import traceback
661        traceback.print_stack()
662        print(f"Failed to find site-packages directory for virtual environment '{venv}'.")
663        print("This may be because you are using a different Python version.")
664        print("Try deleting the following directory and restarting Meerschaum:")
665        print(VIRTENV_RESOURCES_PATH)
666        sys.exit(1)
667
668    return target_path
669
670
671def inside_venv() -> bool:
672    """
673    Determine whether current Python interpreter is running inside a virtual environment.
674    """
675    import sys
676    return (
677        hasattr(sys, 'real_prefix') or (
678            hasattr(sys, 'base_prefix')
679                and sys.base_prefix != sys.prefix
680        )
681    )
682
683
684def get_venvs() -> List[str]:
685    """
686    Return a list of all the virtual environments.
687    """
688    import os
689    from meerschaum.config._paths import VIRTENV_RESOURCES_PATH
690    venvs = []
691    for filename in os.listdir(VIRTENV_RESOURCES_PATH):
692        path = VIRTENV_RESOURCES_PATH / filename
693        if not path.is_dir():
694            continue
695        if not venv_exists(filename):
696            continue
697        venvs.append(filename)
698    return venvs
699
700
701from meerschaum.utils.venv._Venv import Venv