1# Copyright 2014 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import codecs
6import os
7import sys
8import collections
9import StringIO
10
11
12class WithableStringIO(StringIO.StringIO):
13
14  def __enter__(self, *args):
15    return self
16
17  def __exit__(self, *args):
18    pass
19
20
21class FakeFS(object):
22
23  def __init__(self, initial_filenames_and_contents=None):
24    self._file_contents = {}
25    if initial_filenames_and_contents:
26      for k, v in initial_filenames_and_contents.iteritems():
27        self._file_contents[k] = v
28
29    self._bound = False
30    self._real_codecs_open = codecs.open
31    self._real_open = sys.modules['__builtin__'].open
32    self._real_abspath = os.path.abspath
33    self._real_exists = os.path.exists
34    self._real_walk = os.walk
35    self._real_listdir = os.listdir
36
37  def __enter__(self):
38    self.Bind()
39    return self
40
41  def __exit__(self, *args):
42    self.Unbind()
43
44  def Bind(self):
45    assert not self._bound
46    codecs.open = self._FakeCodecsOpen
47    sys.modules['__builtin__'].open = self._FakeOpen
48    os.path.abspath = self._FakeAbspath
49    os.path.exists = self._FakeExists
50    os.walk = self._FakeWalk
51    os.listdir = self._FakeListDir
52    self._bound = True
53
54  def Unbind(self):
55    assert self._bound
56    codecs.open = self._real_codecs_open
57    sys.modules['__builtin__'].open = self._real_open
58    os.path.abspath = self._real_abspath
59    os.path.exists = self._real_exists
60    os.walk = self._real_walk
61    os.listdir = self._real_listdir
62    self._bound = False
63
64  def AddFile(self, path, contents):
65    assert path not in self._file_contents
66    path = os.path.normpath(path)
67    self._file_contents[path] = contents
68
69  def _FakeOpen(self, path, mode=None):
70    if mode is None:
71      mode = 'r'
72    if mode == 'r' or mode == 'rU' or mode == 'rb':
73      if path not in self._file_contents:
74        return self._real_open(path, mode)
75      return WithableStringIO(self._file_contents[path])
76
77    raise NotImplementedError()
78
79  def _FakeCodecsOpen(self, path, mode=None,
80                      encoding=None):  # pylint: disable=unused-argument
81    if mode is None:
82      mode = 'r'
83    if mode == 'r' or mode == 'rU' or mode == 'rb':
84      if path not in self._file_contents:
85        return self._real_open(path, mode)
86      return WithableStringIO(self._file_contents[path])
87
88    raise NotImplementedError()
89
90  def _FakeAbspath(self, path):
91    """Normalize the path and ensure it starts with os.path.sep.
92
93    The tests all assume paths start with things like '/my/project',
94    and this abspath implementaion makes that assumption work correctly
95    on Windows.
96    """
97    normpath = os.path.normpath(path)
98    if not normpath.startswith(os.path.sep):
99      normpath = os.path.sep + normpath
100    return normpath
101
102  def _FakeExists(self, path):
103    if path in self._file_contents:
104      return True
105    return self._real_exists(path)
106
107  def _FakeWalk(self, top):
108    assert os.path.isabs(top)
109    all_filenames = self._file_contents.keys()
110    pending_prefixes = collections.deque()
111    pending_prefixes.append(top)
112    visited_prefixes = set()
113    while len(pending_prefixes):
114      prefix = pending_prefixes.popleft()
115      if prefix in visited_prefixes:
116        continue
117      visited_prefixes.add(prefix)
118      if prefix.endswith(os.path.sep):
119        prefix_with_trailing_sep = prefix
120      else:
121        prefix_with_trailing_sep = prefix + os.path.sep
122
123      dirs = set()
124      files = []
125      for filename in all_filenames:
126        if not filename.startswith(prefix_with_trailing_sep):
127          continue
128        relative_to_prefix = os.path.relpath(filename, prefix)
129
130        dirpart = os.path.dirname(relative_to_prefix)
131        if len(dirpart) == 0:
132          files.append(relative_to_prefix)
133          continue
134        parts = dirpart.split(os.sep)
135        if len(parts) == 0:
136          dirs.add(dirpart)
137        else:
138          pending = os.path.join(prefix, parts[0])
139          dirs.add(parts[0])
140          pending_prefixes.appendleft(pending)
141
142      dirs = sorted(dirs)
143      yield prefix, dirs, files
144
145  def _FakeListDir(self, dirname):
146    raise NotImplementedError()
147