1import abc
2import os
3import tarfile
4
5
6class FileSystemError(Exception):
7    pass
8
9
10class FileSystem(abc.ABC):
11    """Interface for file systems."""
12
13    @abc.abstractmethod
14    def isfile(self, path):
15        """Is this a file?"""
16        pass
17
18    @abc.abstractmethod
19    def isdir(self, path):
20        """Is this a directory?"""
21        pass
22
23    @abc.abstractmethod
24    def read(self, path):
25        """Read a file."""
26        pass
27
28    @abc.abstractmethod
29    def refer_to(self, path):
30        """Get a fully qualified path for the given path."""
31        pass
32
33    def relative_path(self, path):
34        """Return the relative path to `path`.
35
36        If this filesystem has a root directory, and path is within that
37        directory tree, return the relative path; otherwise return None.
38        """
39        return None
40
41
42class StoredFileSystem(FileSystem):
43    """File system based on a file list."""
44
45    def __init__(self, files):
46        self.files = files
47        self.dirs = {os.path.dirname(f) for f in files}
48
49    def isfile(self, path):
50        return path in self.files
51
52    def isdir(self, path):
53        return path in self.dirs
54
55    def read(self, path):
56        return self.files[path]
57
58    def refer_to(self, path):
59        return path
60
61
62class OSFileSystem(FileSystem):
63    """File system that uses an OS file system underneath."""
64
65    def __init__(self, root):
66        assert root is not None
67        self.root = root
68
69    def _join(self, path):
70        return os.path.join(self.root, path)
71
72    def isfile(self, path):
73        assert path is not None
74        return os.path.isfile(self._join(path))
75
76    def isdir(self, path):
77        assert path is not None
78        return os.path.isdir(self._join(path))
79
80    def read(self, path):
81        with open(self._join(path), 'r') as fi:
82            return fi.read()
83
84    def refer_to(self, path):
85        return self._join(path)
86
87    def relative_path(self, path):
88        if path.startswith(self.root):
89            return path[len(self.root) + 1:]
90        return None
91
92
93class RemappingFileSystem(FileSystem, abc.ABC):
94    """File system wrapper that transforms a path before looking it up."""
95
96    def __init__(self, underlying):
97        self.underlying = underlying
98
99    @abc.abstractmethod
100    def map_path(self, path):
101        pass
102
103    def isfile(self, path):
104        return self.underlying.isfile(self.map_path(path))
105
106    def isdir(self, path):
107        return self.underlying.isdir(self.map_path(path))
108
109    def read(self, path):
110        return self.underlying.read(self.map_path(path))
111
112    def refer_to(self, path):
113        return self.underlying.refer_to(self.map_path(path))
114
115
116class ExtensionRemappingFileSystem(RemappingFileSystem):
117    """File system that remaps .py file extensions."""
118
119    def __init__(self, underlying, extension):
120        super(ExtensionRemappingFileSystem, self).__init__(underlying)
121        self.extension = extension
122
123    def map_path(self, path):
124        p, ext = os.path.splitext(path)
125        if ext == '.py':
126            return p + '.' + self.extension
127        return path
128
129
130class PYIFileSystem(ExtensionRemappingFileSystem):
131    """File system that remaps .py file extensions to pyi."""
132
133    def __init__(self, underlying):
134        super(PYIFileSystem, self).__init__(underlying, 'pyi')
135
136
137class TarFileSystem(object):
138    """Filesystem that serves files out of a .tar."""
139
140    def __init__(self, tar):
141        self.tar = tar
142        self.files = list(t.name for t in tar.getmembers() if t.isfile())
143        self.directories = list(t.name for t in tar.getmembers() if t.isdir())
144        self.top_level = {f.split(os.path.sep)[0] for f in self.files}
145
146    def isfile(self, path):
147        return any(os.path.join(top, path) in self.files
148                   for top in self.top_level)
149
150    def isdir(self, path):
151        return any(os.path.join(top, path) in self.files
152                   for top in self.top_level)
153
154    def read(self, path):
155        return self.tar.extractfile(path).read()
156
157    def refer_to(self, path):
158        return 'tar:' + path
159
160    @staticmethod
161    def read_tarfile(archive_filename):
162        tar = tarfile.open(archive_filename)
163        return TarFileSystem(tar)
164
165
166class Path(object):
167    def __init__(self, paths=None):
168        self.paths = paths if paths else []
169
170    def add_path(self, path, kind='os'):
171        if kind == 'os':
172            path = OSFileSystem(path)
173        elif kind == 'pyi':
174            path = PYIFileSystem(OSFileSystem(path))
175        else:
176            raise FileSystemError('Unrecognized filesystem type: ', kind)
177        self.paths.append(path)
178
179    def add_fs(self, fs):
180        assert isinstance(fs, FileSystem), 'Unrecognised filesystem: %r' % fs
181        self.paths.append(fs)
182