← Back to branch summary

~ubuntu-branches/ubuntu/maverick/python3.1/maverick

« back to all changes in this revision

Viewing changes to Lib/importlib/test/source/util.py

  • Committer: Bazaar Package Importer
  • Author(s): Matthias Klose
  • Date: 2009-03-23 00:01:27 UTC
  • Revision ID: james.westby@ubuntu.com-20090323000127-5fstfxju4ufrhthq
Tags: upstream-3.1~a1+20090322
Import upstream version 3.1~a1+20090322

Show diffs side-by-side

added added

removed removed

Lines of Context:
Lib/importlib/test/source/util.py
 
1
from .. import util
 
2
import contextlib
 
3
import functools
 
4
import imp
 
5
import os
 
6
import os.path
 
7
import sys
 
8
import tempfile
 
9
from test import support
 
10
 
 
11
 
 
12
def writes_bytecode(fxn):
 
13
    """Decorator to protect sys.dont_write_bytecode from mutation."""
 
14
    @functools.wraps(fxn)
 
15
    def wrapper(*args, **kwargs):
 
16
        original = sys.dont_write_bytecode
 
17
        sys.dont_write_bytecode = False
 
18
        to_return = fxn(*args, **kwargs)
 
19
        sys.dont_write_bytecode = original
 
20
        return to_return
 
21
    return wrapper
 
22
 
 
23
 
 
24
def writes_bytecode_files(fxn):
 
25
    """Decorator that returns the function if writing bytecode is enabled, else
 
26
    a stub function that accepts anything and simply returns None."""
 
27
    if sys.dont_write_bytecode:
 
28
        return lambda *args, **kwargs: None
 
29
    else:
 
30
        @functools.wraps(fxn)
 
31
        def wrapper(*args, **kwargs):
 
32
            to_return = fxn(*args, **kwargs)
 
33
            sys.dont_write_bytecode = False
 
34
            return to_return
 
35
        return wrapper
 
36
 
 
37
 
 
38
def bytecode_path(source_path):
 
39
    for suffix, _, type_ in imp.get_suffixes():
 
40
        if type_ == imp.PY_COMPILED:
 
41
            bc_suffix = suffix
 
42
            break
 
43
    else:
 
44
        raise ValueError("no bytecode suffix is defined")
 
45
    return os.path.splitext(source_path)[0] + bc_suffix
 
46
 
 
47
 
 
48
@contextlib.contextmanager
 
49
def create_modules(*names):
 
50
    """Temporarily create each named module with an attribute (named 'attr')
 
51
    that contains the name passed into the context manager that caused the
 
52
    creation of the module.
 
53
 
 
54
    All files are created in a temporary directory specified by
 
55
    tempfile.gettempdir(). This directory is inserted at the beginning of
 
56
    sys.path. When the context manager exits all created files (source and
 
57
    bytecode) are explicitly deleted.
 
58
 
 
59
    No magic is performed when creating packages! This means that if you create
 
60
    a module within a package you must also create the package's __init__ as
 
61
    well.
 
62
 
 
63
    """
 
64
    source = 'attr = {0!r}'
 
65
    created_paths = []
 
66
    mapping = {}
 
67
    try:
 
68
        temp_dir = tempfile.gettempdir()
 
69
        mapping['.root'] = temp_dir
 
70
        import_names = set()
 
71
        for name in names:
 
72
            if not name.endswith('__init__'):
 
73
                import_name = name
 
74
            else:
 
75
                import_name = name[:-len('.__init__')]
 
76
            import_names.add(import_name)
 
77
            if import_name in sys.modules:
 
78
                del sys.modules[import_name]
 
79
            name_parts = name.split('.')
 
80
            file_path = temp_dir
 
81
            for directory in name_parts[:-1]:
 
82
                file_path = os.path.join(file_path, directory)
 
83
                if not os.path.exists(file_path):
 
84
                    os.mkdir(file_path)
 
85
                    created_paths.append(file_path)
 
86
            file_path = os.path.join(file_path, name_parts[-1] + '.py')
 
87
            with open(file_path, 'w') as file:
 
88
                file.write(source.format(name))
 
89
            created_paths.append(file_path)
 
90
            mapping[name] = file_path
 
91
        uncache_manager = util.uncache(*import_names)
 
92
        uncache_manager.__enter__()
 
93
        state_manager = util.import_state(path=[temp_dir])
 
94
        state_manager.__enter__()
 
95
        yield mapping
 
96
    finally:
 
97
        state_manager.__exit__(None, None, None)
 
98
        uncache_manager.__exit__(None, None, None)
 
99
        # Reverse the order for path removal to unroll directory creation.
 
100
        for path in reversed(created_paths):
 
101
            if file_path.endswith('.py'):
 
102
                support.unlink(path)
 
103
                support.unlink(path + 'c')
 
104
                support.unlink(path + 'o')
 
105
            else:
 
106
                os.rmdir(path)