1# -*- coding: utf-8 -*- #
2# Copyright 2014 Google LLC. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#    http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Utilities for accessing local package resources."""
17
18from __future__ import absolute_import
19from __future__ import division
20from __future__ import unicode_literals
21
22import imp
23import os
24import pkgutil
25import sys
26
27from googlecloudsdk.core.util import files
28
29import six
30
31
32def _GetPackageName(module_name):
33  """Returns package name for given module name."""
34  last_dot_idx = module_name.rfind('.')
35  if last_dot_idx > 0:
36    return module_name[:last_dot_idx]
37  return ''
38
39
40def GetResource(module_name, resource_name):
41  """Get a resource as a byte string for given resource in same package."""
42  return pkgutil.get_data(_GetPackageName(module_name), resource_name)
43
44
45def GetResourceFromFile(path):
46  """Gets the given resource as a byte string.
47
48  This is similar to GetResource(), but uses file paths instead of module names.
49
50  Args:
51    path: str, filesystem like path to a file/resource.
52
53  Returns:
54    The contents of the resource as a byte string.
55
56  Raises:
57    IOError: if resource is not found under given path.
58  """
59  if os.path.isfile(path):
60    return files.ReadBinaryFileContents(path)
61
62  importer = pkgutil.get_importer(os.path.dirname(path))
63  if hasattr(importer, 'get_data'):
64    return importer.get_data(path)
65
66  raise IOError('File not found {0}'.format(path))
67
68
69def IsImportable(name, path):
70  """Checks if given name can be imported at given path.
71
72  Args:
73    name: str, module name without '.' or suffixes.
74    path: str, filesystem path to location of the module.
75
76  Returns:
77    True, if name is importable.
78  """
79
80  if os.path.isdir(path):
81    if not os.path.isfile(os.path.join(path, '__init__.py')):
82      return path in sys.path
83    name_path = os.path.join(path, name)
84    if os.path.isdir(name_path):
85      # Subdirectory is considered subpackage if it has __init__.py file.
86      return os.path.isfile(os.path.join(name_path, '__init__.py'))
87    return os.path.exists(name_path + '.py')
88
89  try:
90    result = imp.find_module(name, [path])
91    if result:
92      return True
93  except ImportError:
94    pass
95
96  if not hasattr(pkgutil, 'get_importer'):
97    return False
98
99  name_path = name.split('.')
100  importer = pkgutil.get_importer(os.path.join(path, *name_path[:-1]))
101
102  return importer and importer.find_module(name_path[-1])
103
104
105def _GetPathRoot(path):
106  """Returns longest path from sys.path which is prefix of given path."""
107
108  longest_path = ''
109  for p in sys.path:
110    if path.startswith(p) and len(longest_path) < len(p):
111      longest_path = p
112  return longest_path
113
114
115def GetModuleFromPath(name_to_give, module_path):
116  """Loads module at given path under given name.
117
118  Note that it also updates sys.modules with name_to_give.
119
120  Args:
121    name_to_give: str, name to assign to loaded module
122    module_path: str, python path to location of the module, this is either
123        filesystem path or path into egg or zip package
124
125  Returns:
126    Imported module
127
128  Raises:
129    ImportError: if module cannot be imported.
130  """
131  module_dir, module_name = os.path.split(module_path)
132  try:
133    result = imp.find_module(module_name, [module_dir])
134  except ImportError:
135    # imp.find_module does not respects PEP 302 import hooks, and does not work
136    # over package archives. Try pkgutil import hooks.
137    return _GetModuleFromPathViaPkgutil(module_path, name_to_give)
138  else:
139    try:
140      f, file_path, items = result
141      module = imp.load_module(name_to_give, f, file_path, items)
142      if module.__name__ not in sys.modules:
143        # Python 2.6 does not add this to sys.modules. This is to make sure
144        # we get uniform behaviour with 2.7.
145        sys.modules[module.__name__] = module
146      return module
147    finally:
148      if f:
149        f.close()
150
151
152def _GetModuleFromPathViaPkgutil(module_path, name_to_give):
153  """Loads module by using pkgutil.get_importer mechanism."""
154  importer = pkgutil.get_importer(os.path.dirname(module_path))
155  if importer:
156    if hasattr(importer, '_par'):
157      # par zipimporters must have full path from the zip root.
158      # pylint:disable=protected-access
159      module_name = '.'.join(
160          module_path[len(importer._par._zip_filename) + 1:].split(os.sep))
161    else:
162      module_name = os.path.basename(module_path)
163
164    if importer.find_module(module_name):
165      return _LoadModule(importer, module_path, module_name, name_to_give)
166
167  raise ImportError('{0} not found'.format(module_path))
168
169
170def _LoadModule(importer, module_path, module_name, name_to_give):
171  """Loads the module or package under given name."""
172  code = importer.get_code(module_name)
173  module = imp.new_module(name_to_give)
174  package_path_parts = name_to_give.split('.')
175  if importer.is_package(module_name):
176    module.__path__ = [module_path]
177    module.__file__ = os.path.join(module_path, '__init__.pyc')
178  else:
179    package_path_parts.pop()  # Don't treat module as a package.
180    module.__file__ = module_path + '.pyc'
181
182  # Define package if it does not exists.
183  if six.PY2:
184    # This code does not affect the official installations of the cloud sdk.
185    # This function does not work on python 3, but removing this call will
186    # generate runtime warnings when running gcloud as a zip archive. So we keep
187    # this call in python 2 so it can continue to work as intended.
188    imp.load_module('.'.join(package_path_parts), None,
189                    os.path.join(_GetPathRoot(module_path),
190                                 *package_path_parts),
191                    ('', '', imp.PKG_DIRECTORY))
192
193  # pylint: disable=exec-used
194  exec(code, module.__dict__)
195  sys.modules[name_to_give] = module
196  return module
197
198
199def _IterModules(file_list, extra_extensions, prefix=None):
200  """Yields module names from given list of file paths with given prefix."""
201  yielded = set()
202  if extra_extensions is None:
203    extra_extensions = []
204  if prefix is None:
205    prefix = ''
206  for file_path in file_list:
207    if not file_path.startswith(prefix):
208      continue
209
210    file_path_parts = file_path[len(prefix):].split(os.sep)
211
212    if (len(file_path_parts) == 2
213        and file_path_parts[1].startswith('__init__.py')):
214      if file_path_parts[0] not in yielded:
215        yielded.add(file_path_parts[0])
216        yield file_path_parts[0], True
217
218    if len(file_path_parts) != 1:
219      continue
220
221    filename = os.path.basename(file_path_parts[0])
222    modname, ext = os.path.splitext(filename)
223    if modname == '__init__' or (ext != '.py' and ext not in extra_extensions):
224      continue
225
226    to_yield = modname if ext == '.py' else filename
227    if '.' not in modname and to_yield not in yielded:
228      yielded.add(to_yield)
229      yield to_yield, False
230
231
232def _ListPackagesAndFiles(path):
233  """List packages or modules which can be imported at given path."""
234  importables = []
235  for filename in os.listdir(path):
236    if os.path.isfile(os.path.join(path, filename)):
237      importables.append(filename)
238    else:
239      pkg_init_filepath = os.path.join(path, filename, '__init__.py')
240      if os.path.isfile(pkg_init_filepath):
241        importables.append(os.path.join(filename, '__init__.py'))
242  return importables
243
244
245def ListPackage(path, extra_extensions=None):
246  """Returns list of packages and modules in given path.
247
248  Args:
249    path: str, filesystem path
250    extra_extensions: [str], The list of file extra extensions that should be
251      considered modules for the purposes of listing (in addition to .py).
252
253  Returns:
254    tuple([packages], [modules])
255  """
256  iter_modules = []
257  if os.path.isdir(path):
258    iter_modules = _IterModules(_ListPackagesAndFiles(path), extra_extensions)
259  else:
260    importer = pkgutil.get_importer(path)
261    if hasattr(importer, '_files'):
262      # pylint:disable=protected-access
263      iter_modules = _IterModules(
264          importer._files, extra_extensions, importer.prefix)
265    elif hasattr(importer, '_par'):
266      # pylint:disable=protected-access
267      prefix = os.path.join(*importer._prefix.split('.'))
268      iter_modules = _IterModules(
269          importer._par._filename_list, extra_extensions, prefix)
270    elif hasattr(importer, 'ziparchive'):
271      prefix = os.path.join(*importer.prefix.split('.'))
272      # pylint:disable=protected-access
273      iter_modules = _IterModules(
274          importer.ziparchive._files, extra_extensions, prefix)
275  packages, modules = [], []
276  for name, ispkg in iter_modules:
277    if ispkg:
278      packages.append(name)
279    else:
280      modules.append(name)
281  return sorted(packages), sorted(modules)
282
283
284def _IterPrefixFiles(file_list, prefix=None, depth=0):
285  """Returns list of files located at specified prefix dir.
286
287  Args:
288    file_list: list(str), filepaths, usually absolute.
289    prefix: str, filepath prefix, usually proper path itself. Used to filter
290        out files in files_list.
291    depth: int, relative to prefix, of whether to returns files in
292        subdirectories. Depth of 0 would return files in prefix directory.
293
294  Yields:
295    file paths, relative to prefix at given depth or less.
296  """
297  if prefix is None:
298    prefix = ''
299  for file_path in file_list:
300    if not file_path.startswith(prefix):
301      continue
302
303    rel_file_path = file_path[len(prefix):]
304
305    sep_count = depth
306    if rel_file_path.endswith(os.sep):
307      sep_count += 1
308    if rel_file_path.count(os.sep) > sep_count:
309      continue
310    yield rel_file_path
311
312
313def ListPackageResources(path):
314  """Returns list of resources at given path.
315
316  Similar to pkg_resources.resource_listdir.
317
318  Args:
319    path: filesystem like path to a directory/package.
320
321  Returns:
322    list of files/resources at specified path.
323  """
324  if os.path.isdir(path):
325    return [f + os.sep if os.path.isdir(os.path.join(path, f)) else f
326            for f in os.listdir(path)]
327
328  importer = pkgutil.get_importer(path)
329  if hasattr(importer, '_files'):
330    # pylint:disable=protected-access
331    return _IterPrefixFiles(importer._files, importer.prefix, 0)
332
333  if hasattr(importer, '_par'):
334    # pylint:disable=protected-access
335    prefix = os.path.join(*importer._prefix.split('.'))
336    return _IterPrefixFiles(importer._par._filename_list, prefix, 0)
337
338  return []
339