1# -*- coding: utf-8 -*- #
2# Copyright 2017 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"""Helpers to load commands from the filesystem."""
17
18from __future__ import absolute_import
19from __future__ import division
20from __future__ import unicode_literals
21
22import abc
23import os
24import re
25
26import googlecloudsdk
27from googlecloudsdk.calliope import base
28from googlecloudsdk.calliope import command_release_tracks
29from googlecloudsdk.core import exceptions
30from googlecloudsdk.core.util import pkg_resources
31
32from ruamel import yaml
33import six
34
35
36class CommandLoadFailure(Exception):
37  """An exception for when a command or group module cannot be imported."""
38
39  def __init__(self, command, root_exception):
40    self.command = command
41    self.root_exception = root_exception
42    super(CommandLoadFailure, self).__init__(
43        'Problem loading {command}: {issue}.'.format(
44            command=command, issue=six.text_type(root_exception)))
45
46
47class LayoutException(Exception):
48  """An exception for when a command or group .py file has the wrong types."""
49
50
51class ReleaseTrackNotImplementedException(Exception):
52  """An exception for when a command or group does not support a release track.
53  """
54
55
56class YamlCommandTranslator(six.with_metaclass(abc.ABCMeta, object)):
57  """An interface to implement when registering a custom command loader."""
58
59  @abc.abstractmethod
60  def Translate(self, path, command_data):
61    """Translates a yaml command into a calliope command.
62
63    Args:
64      path: [str], A list of group names that got us down to this command group
65        with respect to the CLI itself.  This path should be used for things
66        like error reporting when a specific element in the tree needs to be
67        referenced.
68      command_data: dict, The parsed contents of the command spec from the
69        yaml file that corresponds to the release track being loaded.
70
71    Returns:
72      calliope.base.Command, A command class (not instance) that
73      implements the spec.
74    """
75    pass
76
77
78def FindSubElements(impl_paths, path):
79  """Find all the sub groups and commands under this group.
80
81  Args:
82    impl_paths: [str], A list of file paths to the command implementation for
83      this group.
84    path: [str], A list of group names that got us down to this command group
85      with respect to the CLI itself.  This path should be used for things
86      like error reporting when a specific element in the tree needs to be
87      referenced.
88
89  Raises:
90    CommandLoadFailure: If the command is invalid and cannot be loaded.
91    LayoutException: if there is a command or group with an illegal name.
92
93  Returns:
94    ({str: [str]}, {str: [str]), A tuple of groups and commands found where each
95    item is a mapping from name to a list of paths that implement that command
96    or group. There can be multiple paths because a command or group could be
97    implemented in both python and yaml (for different release tracks).
98  """
99  if len(impl_paths) > 1:
100    raise CommandLoadFailure(
101        '.'.join(path),
102        Exception('Command groups cannot be implemented in yaml'))
103  impl_path = impl_paths[0]
104  groups, commands = pkg_resources.ListPackage(
105      impl_path, extra_extensions=['.yaml'])
106  return (_GenerateElementInfo(impl_path, groups),
107          _GenerateElementInfo(impl_path, commands))
108
109
110def _GenerateElementInfo(impl_path, names):
111  """Generates the data a group needs to load sub elements.
112
113  Args:
114    impl_path: The file path to the command implementation for this group.
115    names: [str], The names of the sub groups or commands found in the group.
116
117  Raises:
118    LayoutException: if there is a command or group with an illegal name.
119
120  Returns:
121    {str: [str], A mapping from name to a list of paths that implement that
122    command or group. There can be multiple paths because a command or group
123    could be implemented in both python and yaml (for different release tracks).
124  """
125  elements = {}
126  for name in names:
127    if re.search('[A-Z]', name):
128      raise LayoutException(
129          'Commands and groups cannot have capital letters: {0}.'.format(name))
130    cli_name = name[:-5] if name.endswith('.yaml') else name
131    sub_path = os.path.join(impl_path, name)
132
133    existing = elements.setdefault(cli_name, [])
134    existing.append(sub_path)
135  return elements
136
137
138def LoadCommonType(impl_paths, path, release_track,
139                   construction_id, is_command, yaml_command_translator=None):
140  """Loads a calliope command or group from a file.
141
142  Args:
143    impl_paths: [str], A list of file paths to the command implementation for
144      this group or command.
145    path: [str], A list of group names that got us down to this command group
146      with respect to the CLI itself.  This path should be used for things
147      like error reporting when a specific element in the tree needs to be
148      referenced.
149    release_track: ReleaseTrack, The release track that we should load.
150    construction_id: str, A unique identifier for the CLILoader that is
151      being constructed.
152    is_command: bool, True if we are loading a command, False to load a group.
153    yaml_command_translator: YamlCommandTranslator, An instance of a translator
154      to use to load the yaml data.
155
156  Raises:
157    CommandLoadFailure: If the command is invalid and cannot be loaded.
158
159  Returns:
160    The base._Common class for the command or group.
161  """
162  implementations = _GetAllImplementations(
163      impl_paths, path, construction_id, is_command, yaml_command_translator)
164  return _ExtractReleaseTrackImplementation(
165      impl_paths[0], release_track, implementations)()
166
167
168def Cache(func):
169  cached_results = {}
170  def ReturnCachedOrCallFunc(*args):
171    try:
172      return cached_results[args]
173    except KeyError:
174      result = func(*args)
175      cached_results[args] = result
176      return result
177  return ReturnCachedOrCallFunc
178
179
180@Cache
181def _SafeLoadYamlFile(path):
182  return yaml.safe_load(pkg_resources.GetResourceFromFile(path))
183
184
185@Cache
186def _CustomLoadYamlFile(path):
187  return CreateYamlLoader(path).load(pkg_resources.GetResourceFromFile(path))
188
189
190def _GetAllImplementations(impl_paths, path, construction_id, is_command,
191                           yaml_command_translator):
192  """Gets all the release track command implementations.
193
194  Can load both python and yaml modules.
195
196  Args:
197    impl_paths: [str], A list of file paths to the command implementation for
198      this group or command.
199    path: [str], A list of group names that got us down to this command group
200      with respect to the CLI itself.  This path should be used for things
201      like error reporting when a specific element in the tree needs to be
202      referenced.
203    construction_id: str, A unique identifier for the CLILoader that is
204      being constructed.
205    is_command: bool, True if we are loading a command, False to load a group.
206    yaml_command_translator: YamlCommandTranslator, An instance of a translator
207      to use to load the yaml data.
208
209  Raises:
210    CommandLoadFailure: If the command is invalid and cannot be loaded.
211
212  Returns:
213    [(func->base._Common, [base.ReleaseTrack])], A list of tuples that can be
214    passed to _ExtractReleaseTrackImplementation. Each item in this list
215    represents a command implementation. The first element is a function that
216    returns the implementation, and the second element is a list of release
217    tracks it is valid for.
218  """
219  implementations = []
220  for impl_file in impl_paths:
221    if impl_file.endswith('.yaml'):
222      if not is_command:
223        raise CommandLoadFailure(
224            '.'.join(path),
225            Exception('Command groups cannot be implemented in yaml'))
226      data = _CustomLoadYamlFile(impl_file)
227      implementations.extend((_ImplementationsFromYaml(
228          path, data, yaml_command_translator)))
229    else:
230      module = _GetModuleFromPath(impl_file, path, construction_id)
231      implementations.extend(_ImplementationsFromModule(
232          module.__file__, list(module.__dict__.values()),
233          is_command=is_command))
234  return implementations
235
236
237def CreateYamlLoader(impl_path):
238  """Creates a custom yaml loader that handles includes from common data.
239
240  Args:
241    impl_path: str, The path to the file we are loading data from.
242
243  Returns:
244    yaml.Loader, A yaml loader to use.
245  """
246  # TODO(b/64147277) Allow for importing from other places.
247  common_file_path = os.path.join(os.path.dirname(impl_path), '__init__.yaml')
248  common_data = None
249  try:
250    common_data = _SafeLoadYamlFile(common_file_path)
251  except IOError:
252    pass
253
254  class Constructor(yaml.Constructor):
255    """A custom yaml constructor.
256
257    It adds 2 different import capabilities. Assuming __init__.yaml has the
258    contents:
259
260    foo:
261      a: b
262      c: d
263
264    baz:
265      - e: f
266      - g: h
267
268    The first uses a custom constructor to insert data into your current file,
269    so:
270
271    bar: !COMMON foo.a
272
273    results in:
274
275    bar: b
276
277    The second mechanism overrides construct_mapping and construct_sequence to
278    post process the data and replace the merge macro with keys from the other
279    file. We can't use the custom constructor for this as well because the
280    merge key type in yaml is processed before custom constructors which makes
281    importing and merging not possible. So:
282
283    bar:
284      _COMMON_: foo
285      i: j
286
287    results in:
288
289    bar:
290      a: b
291      c: d
292      i: j
293
294    This can also be used to merge list contexts, so:
295
296    bar:
297      - _COMMON_baz
298      - i: j
299
300    results in:
301
302    bar:
303      - e: f
304      - g: h
305      - i: j
306
307    You may also use the !REF and _REF_ directives in the same way. Instead of
308    pulling from the common file, they can pull from an arbitrary yaml file
309    somewhere in the googlecloudsdk tree. The syntax looks like:
310
311    bar: !REF googlecloudsdk.foo.bar:a.b.c
312
313    This will load googlecloudsdk/foo/bar.yaml and from that file return the
314    a.b.c nested attribute.
315    """
316
317    INCLUDE_COMMON_MACRO = '!COMMON'
318    MERGE_COMMON_MACRO = '_COMMON_'
319    INCLUDE_REF_MACRO = '!REF'
320    MERGE_REF_MACRO = '_REF_'
321
322    def construct_mapping(self, *args, **kwargs):
323      data = super(Constructor, self).construct_mapping(*args, **kwargs)
324      data = self._ConstructMappingHelper(Constructor.MERGE_COMMON_MACRO,
325                                          self._GetCommonData, data)
326      return self._ConstructMappingHelper(Constructor.MERGE_REF_MACRO,
327                                          self._GetRefData, data)
328
329    def _ConstructMappingHelper(self, macro, source_func, data):
330      attribute_path = data.pop(macro, None)
331      if not attribute_path:
332        return data
333
334      modified_data = {}
335      for path in attribute_path.split(','):
336        modified_data.update(source_func(path))
337      # Add the explicit data last so it can override the imports.
338      modified_data.update(data)
339      return modified_data
340
341    def construct_sequence(self, *args, **kwargs):
342      data = super(Constructor, self).construct_sequence(*args, **kwargs)
343      data = self._ConstructSequenceHelper(Constructor.MERGE_COMMON_MACRO,
344                                           self._GetCommonData, data)
345      return self._ConstructSequenceHelper(Constructor.MERGE_REF_MACRO,
346                                           self._GetRefData, data)
347
348    def _ConstructSequenceHelper(self, macro, source_func, data):
349      new_list = []
350      for i in data:
351        if isinstance(i, six.string_types) and i.startswith(macro):
352          attribute_path = i[len(macro):]
353          for path in attribute_path.split(','):
354            new_list.extend(source_func(path))
355        else:
356          new_list.append(i)
357      return new_list
358
359    def IncludeCommon(self, node):
360      attribute_path = self.construct_scalar(node)
361      return self._GetCommonData(attribute_path)
362
363    def IncludeRef(self, node):
364      attribute_path = self.construct_scalar(node)
365      return self._GetRefData(attribute_path)
366
367    def _GetCommonData(self, attribute_path):
368      if not common_data:
369        raise LayoutException(
370            'Command [{}] references [common command] data but it does not '
371            'exist.'.format(impl_path))
372      return self._GetAttribute(common_data, attribute_path, 'common command')
373
374    def _GetRefData(self, path):
375      """Loads the YAML data from the given reference.
376
377      A YAML reference must refer to a YAML file and an attribute within that
378      file to extract.
379
380      Args:
381        path: str, The path of the YAML file to import. It must be in the
382          form of: package.module:attribute.attribute, where the module path is
383          separated from the sub attributes within the YAML by a ':'.
384
385      Raises:
386        LayoutException: If the given module or attribute cannot be loaded.
387
388      Returns:
389        The referenced YAML data.
390      """
391      root = os.path.dirname(os.path.dirname(googlecloudsdk.__file__))
392      parts = path.split(':')
393      if len(parts) != 2:
394        raise LayoutException(
395            'Invalid Yaml reference: [{}]. References must be in the format: '
396            'path(.path)+:attribute(.attribute)*'.format(path))
397      yaml_path = os.path.join(root, *parts[0].split('.'))
398      yaml_path += '.yaml'
399      try:
400        data = _SafeLoadYamlFile(yaml_path)
401      except IOError as e:
402        raise LayoutException(
403            'Failed to load Yaml reference file [{}]: {}'.format(yaml_path, e))
404
405      return self._GetAttribute(data, parts[1], yaml_path)
406
407    def _GetAttribute(self, data, attribute_path, location):
408      value = data
409      for attribute in attribute_path.split('.'):
410        value = value.get(attribute, None)
411        if not value:
412          raise LayoutException(
413              'Command [{}] references [{}] data attribute [{}] in '
414              'path [{}] but it does not exist.'
415              .format(impl_path, location, attribute, attribute_path))
416      return value
417
418  loader = yaml.YAML()
419  loader.Constructor = Constructor
420  loader.constructor.add_constructor(Constructor.INCLUDE_COMMON_MACRO,
421                                     Constructor.IncludeCommon)
422  loader.constructor.add_constructor(Constructor.INCLUDE_REF_MACRO,
423                                     Constructor.IncludeRef)
424  return loader
425
426
427def _GetModuleFromPath(impl_file, path, construction_id):
428  """Import the module and dig into it to return the namespace we are after.
429
430  Import the module relative to the top level directory.  Then return the
431  actual module corresponding to the last bit of the path.
432
433  Args:
434    impl_file: str, The path to the file this was loaded from (for error
435      reporting).
436    path: [str], A list of group names that got us down to this command group
437      with respect to the CLI itself.  This path should be used for things
438      like error reporting when a specific element in the tree needs to be
439      referenced.
440    construction_id: str, A unique identifier for the CLILoader that is
441      being constructed.
442
443  Returns:
444    The imported module.
445  """
446  # Make sure this module name never collides with any real module name.
447  # Use the CLI naming path, so values are always unique.
448  name_to_give = '__calliope__command__.{construction_id}.{name}'.format(
449      construction_id=construction_id,
450      name='.'.join(path).replace('-', '_'))
451  try:
452    return pkg_resources.GetModuleFromPath(name_to_give, impl_file)
453  # pylint:disable=broad-except, We really do want to catch everything here,
454  # because if any exceptions make it through for any single command or group
455  # file, the whole CLI will not work. Instead, just log whatever it is.
456  except Exception as e:
457    exceptions.reraise(CommandLoadFailure('.'.join(path), e))
458
459
460def _ImplementationsFromModule(mod_file, module_attributes, is_command):
461  """Gets all the release track command implementations from the module.
462
463  Args:
464    mod_file: str, The __file__ attribute of the module resulting from
465      importing the file containing a command.
466    module_attributes: The __dict__.values() of the module.
467    is_command: bool, True if we are loading a command, False to load a group.
468
469  Raises:
470    LayoutException: If there is not exactly one type inheriting CommonBase.
471
472  Returns:
473    [(func->base._Common, [base.ReleaseTrack])], A list of tuples that can be
474    passed to _ExtractReleaseTrackImplementation. Each item in this list
475    represents a command implementation. The first element is a function that
476    returns the implementation, and the second element is a list of release
477    tracks it is valid for.
478  """
479  commands = []
480  groups = []
481
482  # Collect all the registered groups and commands.
483  for command_or_group in module_attributes:
484    if getattr(command_or_group, 'IS_COMMAND', False):
485      commands.append(command_or_group)
486    elif getattr(command_or_group, 'IS_COMMAND_GROUP', False):
487      groups.append(command_or_group)
488
489  if is_command:
490    if groups:
491      # Ensure that there are no groups if we are expecting a command.
492      raise LayoutException(
493          'You cannot define groups [{0}] in a command file: [{1}]'
494          .format(', '.join([g.__name__ for g in groups]), mod_file))
495    if not commands:
496      # Make sure we found a command.
497      raise LayoutException('No commands defined in file: [{0}]'.format(
498          mod_file))
499    commands_or_groups = commands
500  else:
501    # Ensure that there are no commands if we are expecting a group.
502    if commands:
503      raise LayoutException(
504          'You cannot define commands [{0}] in a command group file: [{1}]'
505          .format(', '.join([c.__name__ for c in commands]), mod_file))
506    if not groups:
507      # Make sure we found a group.
508      raise LayoutException('No command groups defined in file: [{0}]'.format(
509          mod_file))
510    commands_or_groups = groups
511
512  # pylint:disable=undefined-loop-variable, Linter is just wrong here.
513  # We need to use a default param on the lambda so that it captures the value
514  # of the variable at the time in the loop or else the closure will just have
515  # the last value that was iterated on.
516  return [(lambda c=c: c, c.ValidReleaseTracks()) for c in commands_or_groups]
517
518
519def _ImplementationsFromYaml(path, data, yaml_command_translator):
520  """Gets all the release track command implementations from the yaml file.
521
522  Args:
523    path: [str], A list of group names that got us down to this command group
524      with respect to the CLI itself.  This path should be used for things
525      like error reporting when a specific element in the tree needs to be
526      referenced.
527    data: dict, The loaded yaml data.
528    yaml_command_translator: YamlCommandTranslator, An instance of a translator
529      to use to load the yaml data.
530
531  Raises:
532    CommandLoadFailure: If the command is invalid and cannot be loaded.
533
534  Returns:
535    [(func->base._Common, [base.ReleaseTrack])], A list of tuples that can be
536    passed to _ExtractReleaseTrackImplementation. Each item in this list
537    represents a command implementation. The first element is a function that
538    returns the implementation, and the second element is a list of release
539    tracks it is valid for.
540  """
541  if not yaml_command_translator:
542    raise CommandLoadFailure(
543        '.'.join(path),
544        Exception('No yaml command translator has been registered'))
545
546  # pylint:disable=undefined-loop-variable, Linter is just wrong here.
547  # We need to use a default param on the lambda so that it captures the value
548  # of the variable at the time in the loop or else the closure will just have
549  # the last value that was iterated on.
550  implementations = [
551      (lambda i=i: yaml_command_translator.Translate(path, i),
552       {base.ReleaseTrack.FromId(t) for t in i.get('release_tracks', [])})
553      for i in command_release_tracks.SeparateDeclarativeCommandTracks(data)]
554  return implementations
555
556
557def _ExtractReleaseTrackImplementation(
558    impl_file, expected_track, implementations):
559  """Validates and extracts the correct implementation of the command or group.
560
561  Args:
562    impl_file: str, The path to the file this was loaded from (for error
563      reporting).
564    expected_track: base.ReleaseTrack, The release track we are trying to load.
565    implementations: [(func->base._Common, [base.ReleaseTrack])], A list of
566    tuples where each item in this list represents a command implementation. The
567    first element is a function that returns the implementation, and the second
568    element is a list of release tracks it is valid for.
569
570  Raises:
571    LayoutException: If there is not exactly one type inheriting
572        CommonBase.
573    ReleaseTrackNotImplementedException: If there is no command or group
574      implementation for the request release track.
575
576  Returns:
577    object, The single implementation that matches the expected release track.
578  """
579  # We found a single thing, if it's valid for this track, return it.
580  if len(implementations) == 1:
581    impl, valid_tracks = implementations[0]
582    # If there is a single thing defined, and it does not declare any valid
583    # tracks, just assume it is enabled for all tracks that it's parent is.
584    if not valid_tracks or expected_track in valid_tracks:
585      return impl
586    raise ReleaseTrackNotImplementedException(
587        'No implementation for release track [{0}] for element: [{1}]'
588        .format(expected_track.id, impl_file))
589
590  # There was more than one thing found, make sure there are no conflicts.
591  implemented_release_tracks = set()
592  for impl, valid_tracks in implementations:
593    # When there are multiple definitions, they need to explicitly register
594    # their track to keep things sane.
595    if not valid_tracks:
596      raise LayoutException(
597          'Multiple implementations defined for element: [{0}]. Each must '
598          'explicitly declare valid release tracks.'.format(impl_file))
599    # Make sure no two classes define the same track.
600    duplicates = implemented_release_tracks & valid_tracks
601    if duplicates:
602      raise LayoutException(
603          'Multiple definitions for release tracks [{0}] for element: [{1}]'
604          .format(', '.join([six.text_type(d) for d in duplicates]), impl_file))
605    implemented_release_tracks |= valid_tracks
606
607  valid_commands_or_groups = [impl for impl, valid_tracks in implementations
608                              if expected_track in valid_tracks]
609  # We know there is at most 1 because of the above check.
610  if len(valid_commands_or_groups) != 1:
611    raise ReleaseTrackNotImplementedException(
612        'No implementation for release track [{0}] for element: [{1}]'
613        .format(expected_track.id, impl_file))
614
615  return valid_commands_or_groups[0]
616