-
-
Notifications
You must be signed in to change notification settings - Fork 606
/
rules.py
126 lines (113 loc) · 4.61 KB
/
rules.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).
from pants.backend.codegen.protobuf.protoc import Protoc
from pants.backend.codegen.protobuf.target_types import ProtobufSources
from pants.backend.python.target_types import PythonSources
from pants.core.util_rules.determine_source_files import AllSourceFilesRequest, SourceFiles
from pants.core.util_rules.external_tool import DownloadedExternalTool, ExternalToolRequest
from pants.core.util_rules.strip_source_roots import representative_path_from_address
from pants.engine.addresses import Addresses
from pants.engine.fs import AddPrefix, Digest, MergeDigests, RemovePrefix, Snapshot
from pants.engine.platform import Platform
from pants.engine.process import Process, ProcessResult
from pants.engine.rules import register_rules, rule
from pants.engine.selectors import Get, MultiGet
from pants.engine.target import GeneratedSources, GenerateSourcesRequest, Sources, TransitiveTargets
from pants.engine.unions import UnionRule
from pants.source.source_root import SourceRoot, SourceRootRequest
class GeneratePythonFromProtobufRequest(GenerateSourcesRequest):
input = ProtobufSources
output = PythonSources
@rule(desc="Generate Python from Protobuf")
async def generate_python_from_protobuf(
request: GeneratePythonFromProtobufRequest, protoc: Protoc
) -> GeneratedSources:
download_protoc_request = Get(
DownloadedExternalTool, ExternalToolRequest, protoc.get_request(Platform.current)
)
output_dir = "_generated_files"
# TODO(#9650): replace this with a proper intrinsic to create empty directories.
create_output_dir_request = Get(
ProcessResult,
Process(
("/bin/mkdir", output_dir),
description=f"Create the directory {output_dir}",
output_directories=(output_dir,),
),
)
# Protoc needs all transitive dependencies on `protobuf_libraries` to work properly. It won't
# actually generate those dependencies; it only needs to look at their .proto files to work
# with imports.
transitive_targets = await Get(TransitiveTargets, Addresses([request.protocol_target.address]))
all_sources_request = Get(
SourceFiles,
AllSourceFilesRequest(
(tgt.get(Sources) for tgt in transitive_targets.closure),
for_sources_types=(ProtobufSources,),
# NB: By stripping the source roots, we avoid having to set the value `--proto_path`
# for Protobuf imports to be discoverable.
strip_source_roots=True,
),
)
stripped_target_sources_request = Get(
SourceFiles,
AllSourceFilesRequest([request.protocol_target[ProtobufSources]], strip_source_roots=True),
)
(
downloaded_protoc_binary,
create_output_dir_result,
all_sources,
stripped_target_sources,
) = await MultiGet(
download_protoc_request,
create_output_dir_request,
all_sources_request,
stripped_target_sources_request,
)
input_digest = await Get(
Digest,
MergeDigests(
(
all_sources.snapshot.digest,
downloaded_protoc_binary.digest,
create_output_dir_result.output_digest,
)
),
)
result = await Get(
ProcessResult,
Process(
(
downloaded_protoc_binary.exe,
"--python_out",
output_dir,
*stripped_target_sources.snapshot.files,
),
input_digest=input_digest,
description=f"Generating Python sources from {request.protocol_target.address}.",
output_directories=(output_dir,),
),
)
# We must do some path manipulation on the output digest for it to look like normal sources,
# including adding back the original source root.
normalized_digest, source_root = await MultiGet(
Get(Digest, RemovePrefix(result.output_digest, output_dir)),
Get(
SourceRoot,
SourceRootRequest,
SourceRootRequest.for_file(
representative_path_from_address(request.protocol_target.address)
),
),
)
source_root_restored = (
await Get(Snapshot, AddPrefix(normalized_digest, source_root.path))
if source_root.path != "."
else await Get(Snapshot, Digest, normalized_digest)
)
return GeneratedSources(source_root_restored)
def rules():
return [
*register_rules(),
UnionRule(GenerateSourcesRequest, GeneratePythonFromProtobufRequest),
]