237 lines
8.1 KiB
Python
237 lines
8.1 KiB
Python
# -*- coding: utf-8 -*- #
|
|
# Copyright 2023 Google LLC. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Provides a parser for --container arguments."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import collections
|
|
from collections.abc import Sequence
|
|
import re
|
|
from typing import Any
|
|
|
|
from googlecloudsdk.calliope import base as calliope_base
|
|
from googlecloudsdk.calliope import cli
|
|
from googlecloudsdk.calliope import parser_arguments
|
|
from googlecloudsdk.calliope import parser_errors
|
|
from googlecloudsdk.calliope import parser_extensions
|
|
from googlecloudsdk.command_lib.run import flags
|
|
|
|
|
|
def AddContainerFlags(
|
|
parser: parser_arguments.ArgumentInterceptor,
|
|
container_arg_group: calliope_base.ArgumentGroup,
|
|
release_track=calliope_base.ReleaseTrack.GA,
|
|
):
|
|
"""AddContainerFlags updates parser to add --container arg parsing.
|
|
|
|
Args:
|
|
parser: The parser to patch.
|
|
container_arg_group: Arguments that can be specified per-container.
|
|
release_track: The release track of the command.
|
|
"""
|
|
flags.ContainerFlag().AddToParser(parser)
|
|
container_arg_group.AddToParser(parser)
|
|
container_parser = ContainerParser(
|
|
parser.parser, container_arg_group, release_track
|
|
)
|
|
parser.parser.parse_known_args = container_parser.ParseKnownArgs
|
|
|
|
|
|
class ContainerParser(object):
|
|
"""ContainerParser adds custom container parsing behavior to ArgumentParser."""
|
|
|
|
_CONTAINER_FLAG_NAME = '--container'
|
|
_PRESET_FLAG_NAME = '--preset'
|
|
_FLAG_PATTERN = r'^(--[^:=]+)[:=]?.*'
|
|
|
|
def __init__(
|
|
self,
|
|
parser: parser_extensions.ArgumentParser,
|
|
container_arg_group: calliope_base.ArgumentGroup,
|
|
release_track: calliope_base.ReleaseTrack,
|
|
):
|
|
"""ContainerParser constructor.
|
|
|
|
Args:
|
|
parser: The original command's parser. Used to parse non-container args.
|
|
container_arg_group: Arguments to add to per-container parsers.
|
|
release_track: The release track of the command.
|
|
"""
|
|
self._parse_known_args = parser.parse_known_args
|
|
self._prog = parser.prog
|
|
self._calliope_command = parser._calliope_command
|
|
self._container_arg_group = container_arg_group
|
|
self._release_track = release_track
|
|
|
|
def _GetContainerFlags(self) -> frozenset[str]:
|
|
"""_GetContainerFlags returns the configured set of per-container flags."""
|
|
args = [self._container_arg_group]
|
|
flag_names = []
|
|
while args:
|
|
arg = args.pop()
|
|
if isinstance(arg, calliope_base.ArgumentGroup):
|
|
args.extend(arg.arguments)
|
|
else:
|
|
flag_names.append(arg.name)
|
|
return frozenset(flag_names)
|
|
|
|
def _NewContainerParser(self) -> parser_extensions.ArgumentParser:
|
|
"""_NewContainerParser creates a new parser for parsing container args."""
|
|
parser = parser_extensions.ArgumentParser(
|
|
add_help=False,
|
|
prog=self._prog,
|
|
calliope_command=self._calliope_command,
|
|
)
|
|
|
|
ai = parser_arguments.ArgumentInterceptor(
|
|
parser=parser,
|
|
is_global=False,
|
|
cli_generator=None,
|
|
allow_positional=True,
|
|
)
|
|
|
|
self._container_arg_group.AddToParser(ai)
|
|
cli.FLAG_INTERNAL_FLAG_FILE_LINE.AddToParser(ai)
|
|
return parser
|
|
|
|
def _CheckForContainerFlags(self, namespace: parser_extensions.Namespace):
|
|
"""_CheckForContainerFlags checks that no container flags were specified.
|
|
|
|
Args:
|
|
namespace: The namespace to check.
|
|
"""
|
|
container_flags = self._GetContainerFlags().intersection(
|
|
namespace.GetSpecifiedArgNames()
|
|
)
|
|
if container_flags:
|
|
raise parser_errors.ArgumentError(
|
|
'When --container is specified {flags} must be specified after'
|
|
' --container.',
|
|
flags=', '.join(container_flags),
|
|
)
|
|
|
|
def _IsFlagArg(self, arg: Any):
|
|
return isinstance(arg, str) and arg.startswith('--')
|
|
|
|
def _ExtractFlag(self, arg: str):
|
|
flag = re.match(self._FLAG_PATTERN, arg)
|
|
if flag:
|
|
return flag.group(1)
|
|
return None
|
|
|
|
def ParseKnownArgs(
|
|
self,
|
|
args: Sequence[Any],
|
|
namespace: parser_extensions.Namespace,
|
|
) -> tuple[parser_extensions.Namespace, Sequence[Any]]:
|
|
"""Performs custom --container arg parsing.
|
|
|
|
Groups arguments after each --container flag to be parsed into that
|
|
container's namespace. For each container a new parser is used to parse that
|
|
container's flags into fresh namespace and those namespaces are stored as a
|
|
dict in namespace.containers. Remaining args are parsed by the orignal
|
|
parser's parse_known_args method.
|
|
|
|
Args:
|
|
args: The arguments to parse.
|
|
namespace: The namespace to store parsed args in.
|
|
|
|
Returns:
|
|
A tuple containing the updated namespace and a list of unknown args.
|
|
"""
|
|
remaining = []
|
|
containers = collections.defaultdict(list)
|
|
current = remaining
|
|
i = 0
|
|
while i < len(args):
|
|
value = args[i]
|
|
i += 1
|
|
if value == self._CONTAINER_FLAG_NAME:
|
|
if i >= len(args):
|
|
remaining.append(value)
|
|
else:
|
|
current = containers[args[i]]
|
|
i += 1
|
|
elif isinstance(value, str) and value.startswith(
|
|
self._CONTAINER_FLAG_NAME + '='
|
|
):
|
|
current = containers[value.split(sep='=', maxsplit=1)[1]]
|
|
# Add container for preset to container dict to capture container specific
|
|
# flags for the placeholder container.
|
|
# TODO(b/436350694): Change this to use preset metadata and rework this
|
|
# branch to use regex similar to how the base parsers work.
|
|
# For case "--preset=ollama"
|
|
elif (
|
|
isinstance(value, str)
|
|
and value.startswith(self._PRESET_FLAG_NAME + '=')
|
|
):
|
|
preset_arg = value.split(sep='=', maxsplit=1)[1]
|
|
preset_name = preset_arg.split(':')[0]
|
|
if preset_name in flags.INGRESS_CONTAINER_PRESETS:
|
|
current = containers[preset_name]
|
|
remaining.append(value)
|
|
# For case "--preset ollama"
|
|
elif (
|
|
isinstance(value, str)
|
|
and value.startswith(self._PRESET_FLAG_NAME)
|
|
and i < len(args)
|
|
):
|
|
preset_arg = args[i]
|
|
preset_name = preset_arg.split(':')[0]
|
|
if preset_name in flags.INGRESS_CONTAINER_PRESETS:
|
|
current = containers[args[i]]
|
|
remaining.append(value)
|
|
remaining.append(args[i])
|
|
i += 1
|
|
elif value == '--':
|
|
remaining.append(value)
|
|
remaining.extend(args[i:])
|
|
break
|
|
# For any flags not in the container flag list that come after a container
|
|
# flag, add the flag and all arguments to 'remaining' list until the next
|
|
# flag is encountered.
|
|
elif (
|
|
self._release_track
|
|
in [calliope_base.ReleaseTrack.BETA, calliope_base.ReleaseTrack.ALPHA]
|
|
and containers
|
|
and self._IsFlagArg(value)
|
|
and self._ExtractFlag(value) not in self._GetContainerFlags()
|
|
):
|
|
remaining.append(value)
|
|
while i < len(args) and not self._IsFlagArg(args[i]):
|
|
remaining.append(args[i])
|
|
i += 1
|
|
else:
|
|
current.append(value)
|
|
|
|
if not containers:
|
|
return self._parse_known_args(args=remaining, namespace=namespace)
|
|
|
|
namespace.containers = {}
|
|
# pylint: disable=protected-access
|
|
namespace._specified_args['containers'] = self._CONTAINER_FLAG_NAME
|
|
for container_name, container_args in containers.items():
|
|
container_namespace = parser_extensions.Namespace()
|
|
container_namespace = self._NewContainerParser().parse_args(
|
|
args=container_args, namespace=container_namespace
|
|
)
|
|
namespace.containers[container_name] = container_namespace
|
|
|
|
namespace, unknown_args = self._parse_known_args(
|
|
args=remaining, namespace=namespace
|
|
)
|
|
self._CheckForContainerFlags(namespace)
|
|
return namespace, unknown_args
|