Mirror of https://github.com/roostorg/osprey
github.com/roostorg/osprey
cli: small tweaks to get `push-rules` working (#48)
+14
-50
+2
osprey_worker/pyproject.toml
+2
osprey_worker/pyproject.toml
+3
-3
osprey_worker/src/osprey/worker/lib/cli.py
+3
-3
osprey_worker/src/osprey/worker/lib/cli.py
···
30
30
labels,
31
31
stored_execution_result,
32
32
)
33
-
from osprey.worker.lib.utils.click_utils import EnumChoicePb2 # noqa: E402
33
+
from osprey.worker.lib.utils.click_utils import EnumChoice # noqa: E402
34
34
35
35
36
36
@click.group()
···
204
204
@click.argument('entity_type')
205
205
@click.argument('entity_id')
206
206
@click.argument('label_name')
207
-
@click.argument('label_status', type=EnumChoicePb2(LabelStatus))
207
+
@click.argument('label_status', type=EnumChoice(LabelStatus))
208
208
@click.option(
209
209
'--reason',
210
210
help=(
···
267
267
@click.argument('entity_type')
268
268
@click.argument('entity_ids_file_path')
269
269
@click.argument('label_name')
270
-
@click.argument('label_status', type=EnumChoicePb2(LabelStatus))
270
+
@click.argument('label_status', type=EnumChoice(LabelStatus))
271
271
@click.option(
272
272
'--reason',
273
273
help=(
+8
-4
osprey_worker/src/osprey/worker/lib/sources_publisher.py
+8
-4
osprey_worker/src/osprey/worker/lib/sources_publisher.py
···
52
52
53
53
Returns False if rules failed validation, True otherwise.
54
54
"""
55
+
# bootstrap validators before instance_with_additional_validators copies the validator registry
56
+
udf_registry, _ = bootstrap_udfs()
57
+
bootstrap_ast_validators()
58
+
55
59
lib_validator_registry = ValidatorRegistry.instance_with_additional_validators(
56
60
get_config_registry().get_validator()
57
61
)
58
-
udf_registry, _ = bootstrap_udfs()
59
-
bootstrap_ast_validators()
60
62
61
63
try:
62
64
validated_sources = validate_sources(
···
102
104
suppress_warnings: bool = False,
103
105
) -> bool:
104
106
"""Create a mapping of Osprey Rules/UDFs/variables to their dependencies and upload to BigQuery."""
107
+
# bootstrap validators before instance_with_additional_validators copies the validator registry
108
+
udf_registry, _ = bootstrap_udfs()
109
+
bootstrap_ast_validators()
110
+
105
111
lib_validator_registry = ValidatorRegistry.instance_with_additional_validators(
106
112
get_config_registry().get_validator()
107
113
)
108
-
udf_registry, _ = bootstrap_udfs()
109
-
bootstrap_ast_validators()
110
114
111
115
try:
112
116
validated_sources = validate_sources(
+1
-43
osprey_worker/src/osprey/worker/lib/utils/click_utils.py
+1
-43
osprey_worker/src/osprey/worker/lib/utils/click_utils.py
···
1
1
from enum import Enum
2
-
from typing import Any, Dict, Generic, Optional, Type, TypeVar, Union
2
+
from typing import Dict, Generic, Optional, Type, TypeVar, Union
3
3
4
4
import click
5
5
from click import Context, Parameter
···
34
34
# This is to support when click is given a default that is a variant. If we have a value that is our variant,
35
35
# we can just return it without trying to coerce.
36
36
if isinstance(value, self.enum):
37
-
return value
38
-
39
-
assert isinstance(value, str)
40
-
try:
41
-
return self.choice_map[_to_kebab_case(value)]
42
-
except KeyError:
43
-
self.fail('invalid choice: %s. (choose from %s)' % (value, ', '.join(self.choices)), param, ctx)
44
-
45
-
return None
46
-
47
-
48
-
class EnumChoicePb2(click.Choice):
49
-
"""For use with `@click.option(...)` in order to support for choosing a single value of a protobuf enum.
50
-
51
-
Example usage:
52
-
```
53
-
from some_pb2 import SomeEnum
54
-
55
-
@click.option('--choice', type=EnumChoicePb2(SomeEnum), default=SomeEnum.SOME_VALUE)
56
-
def func(choice):
57
-
click.echo(f"Your choice was: {choice}")
58
-
```
59
-
"""
60
-
61
-
choice_map: Dict[str, int]
62
-
enum_class: Any
63
-
64
-
def __init__(self, enum_class: Any):
65
-
self.enum_class = enum_class
66
-
# Build choice map from enum descriptor
67
-
self.choice_map = {}
68
-
69
-
# Get the enum descriptor and build choices from value names
70
-
for value_descriptor in enum_class.DESCRIPTOR.values:
71
-
kebab_name = _to_kebab_case(value_descriptor.name)
72
-
self.choice_map[kebab_name] = value_descriptor.number
73
-
74
-
super().__init__(choices=list(self.choice_map.keys()))
75
-
76
-
def convert(self, value: Union[str, int], param: Optional[Parameter], ctx: Optional[Context]) -> Optional[int]:
77
-
# If we already have an int (protobuf enum value), return it
78
-
if isinstance(value, int):
79
37
return value
80
38
81
39
assert isinstance(value, str)