Extract rule handling logic into separate functions

Summary: extraction_movie

Reviewed By: stroxler

Differential Revision: D88773202

fbshipit-source-id: ed92e1e2c8aded05f84d7983231f82199f1ad637

Author: grievejia

prelude/python/sourcedb/pyrefly.bxl

@@ -73,11 +73,11 @@ def _clean_module(base_module: str, module: str) -> str:
 # Set common fields for a target entry.
 # NOTE: This function doesn't populate `srcs`. Instead return an empty `srcs` dict for downstream logic to populate.
 def _set_entry_common_fields(
-        entry: dict[str, typing.Any],
+        entry: dict,
         attrs: bxl.LazyAttrs,
         fs: bxl.Filesystem,
         target: bxl.ConfiguredTargetNode,
-        known_targets: dict) -> dict[str, list[str | bxl.EnsuredArtifact]]:
+        known_targets: dict[TargetLabel, bxl.ConfiguredTargetNode]) -> dict[str, list[str | bxl.EnsuredArtifact]]:
     srcs = {}
     deps = []
     entry["srcs"] = srcs
@@ -102,13 +102,136 @@ def _set_entry_common_fields(
 
     return srcs
 
+def _handle_alias(
+        entry: dict,
+        attrs: bxl.LazyAttrs,
+        target: bxl.ConfiguredTargetNode) -> None:
+    actual = attrs.get("actual")
+    if actual != None and actual.value() != None:
+        entry["alias"] = actual.value().raw_target()
+    else:
+        fail("Alias has no actual target".format(target.label))
+
+def _handle_genrule(
+        entry: dict,
+        attrs: bxl.LazyAttrs,
+        fs: bxl.Filesystem,
+        target: bxl.ConfiguredTargetNode,
+        known_targets: dict[TargetLabel, bxl.ConfiguredTargetNode],
+        to_remove: dict[str, None],
+        to_build: list,
+        to_insert: list,
+        genrule_to_insert: list) -> None:
+    labels = attrs.get("labels")
+    if "custom_rule_internal" in labels.value():
+        # skip internal targets, since they will be materialized as part
+        # of their including target
+        to_remove[str(target.label.raw_target())] = None
+        return
+
+    srcs = _set_entry_common_fields(entry, attrs, fs, target, known_targets)
+    to_build.append(target)
+
+    # Only one of these will have a value, but get both of them now since
+    # it's easy.
+    out = attrs.get("out")
+    outs = attrs.get("outs")
+    if out != None and out.value() != None:
+        items = {out.value(): out.value()}
+    elif outs != None and outs.value() != None:
+        items = outs.value()
+    else:
+        fail(target)
+
+    if len(items) == 1:
+        module_path = list(items.keys())[0]
+        for suffix in _MODULE_SUFFIXES:
+            if module_path.endswith(suffix):
+                to_build.append(target)
+                to_insert.append((srcs, _clean_module("", module_path), target.label.with_sub_target()))
+                break
+    else:
+        # we need to handle building the full genrule
+        genrule_to_insert.append((target, entry))
+        for (module_path, src_paths) in items.items():
+            if len(src_paths) < 1:
+                fail("Target src has no output path: {}".format(target.label))
+            for suffix in _MODULE_SUFFIXES:
+                if module_path.endswith(suffix):
+                    _append_or_default_if_module(srcs, _clean_module("", module_path), src_paths[0])
+                    break
+
+def _handle_python_library(
+        entry: dict,
+        attrs: bxl.LazyAttrs,
+        fs: bxl.Filesystem,
+        target: bxl.ConfiguredTargetNode,
+        known_targets: dict[TargetLabel, bxl.ConfiguredTargetNode],
+        to_remove: dict[str, None],
+        to_build: list,
+        to_insert: list) -> None:
+    base_module = attrs.get("base_module")
+    if base_module != None and base_module.value() != None:
+        base_module = base_module.value()
+    else:
+        base_module = target.label.package
+
+    srcs = _set_entry_common_fields(entry, attrs, fs, target, known_targets)
+
+    # add type stubs first, so that they're preferred by Pyrefly
+    type_stubs = attrs.get("type_stubs")
+    type_stubs = {} if type_stubs == None else type_stubs.value()
+    _get_and_materialize_sources(
+        type_stubs,
+        to_remove,
+        base_module,
+        srcs,
+        to_build,
+        to_insert,
+        fs,
+    )
+
+    # python_binary might be the root target, which doesn't have `srcs`.
+    attr_srcs = attrs.get("srcs")
+    attr_srcs = {} if attr_srcs == None else attr_srcs.value()
+    _get_and_materialize_sources(
+        attr_srcs,
+        to_remove,
+        base_module,
+        srcs,
+        to_build,
+        to_insert,
+        fs,
+    )
+
+def _build_generated_files(
+        ctx: bxl.Context,
+        to_build: list,
+        to_insert: list,
+        genrule_to_insert: list) -> None:
+    if not to_build:
+        return
+
+    built = ctx.build(to_build)
+    ensured = ctx.output.ensure_multiple(built)
+
+    # handle individual generated files
+    for srcs, module_name, target in to_insert:
+        _append_or_default_if_module(srcs, module_name, ensured[target][0].rel_path())
+
+    # handle genrules
+    for target, entry in genrule_to_insert:
+        label = target.label.with_sub_target()
+        relative_to = ensured[label][0].rel_path()
+        entry["relative_to"] = relative_to
+
 def _get_and_materialize_sources(
-        to_add: list[typing.Any] | dict[str, typing.Any],
-        to_remove: dict[str, typing.Any],
+        to_add: list | dict,
+        to_remove: dict[str, None],
         base_module: str,
         srcs: dict[str, list[str | bxl.EnsuredArtifact]],
-        to_build: list[typing.Any],
-        to_insert: list[typing.Any],
+        to_build: list,
+        to_insert: list,
         fs: bxl.Filesystem):
     if type(to_add) == type([]):
         for v in to_add:
@@ -181,101 +304,13 @@ def _main(ctx: bxl.Context) -> None:
         entries[str(x.label.raw_target())] = entry
 
         if x.rule_type == "prelude//rules.bzl:alias":
-            actual = attrs.get("actual")
-            if actual != None and actual.value() != None:
-                entry["alias"] = actual.value().raw_target()
-            else:
-                fail("Alias has no actual target".format(x.label))
-
+            _handle_alias(entry, attrs, x)
         elif x.rule_type == "prelude//rules.bzl:genrule":
-            labels = attrs.get("labels")
-            if "custom_rule_internal" in labels.value():
-                # skip internal targets, since they will be materialized as part
-                # of their including target
-                to_remove[str(x.label.raw_target())] = None
-                continue
-
-            srcs = _set_entry_common_fields(entry, attrs, fs, x, known_targets)
-            to_build.append(x)
-
-            # Only one of these will have a value, but get both of them now since
-            # it's easy.
-            out = attrs.get("out")
-            outs = attrs.get("outs")
-            if out != None and out.value() != None:
-                items = {out.value(): out.value()}
-            elif outs != None and outs.value() != None:
-                items = outs.value()
-            else:
-                fail(x)
-
-            if len(items) == 1:
-                module_path = list(items.keys())[0]
-                for suffix in _MODULE_SUFFIXES:
-                    if module_path.endswith(suffix):
-                        to_build.append(x)
-                        to_insert.append((srcs, _clean_module("", module_path), x.label.with_sub_target()))
-                        break
-            else:
-                # we need to handle building the full genrule
-                genrule_to_insert.append((x, entry))
-                for (module_path, src_paths) in items.items():
-                    if len(src_paths) < 1:
-                        fail("Target src has no output path: {}".format(x.label))
-                    for suffix in _MODULE_SUFFIXES:
-                        if module_path.endswith(suffix):
-                            _append_or_default_if_module(srcs, _clean_module("", module_path), src_paths[0])
-                            break
-
+            _handle_genrule(entry, attrs, fs, x, known_targets, to_remove, to_build, to_insert, genrule_to_insert)
         else:
-            # python_library and other rules
-            base_module = attrs.get("base_module")
-            if base_module != None and base_module.value() != None:
-                base_module = base_module.value()
-            else:
-                base_module = x.label.package
-
-            srcs = _set_entry_common_fields(entry, attrs, fs, x, known_targets)
-
-            # add type stubs first, so that they're preferred by Pyrefly
-            type_stubs = attrs.get("type_stubs")
-            type_stubs = {} if type_stubs == None else type_stubs.value()
-            _get_and_materialize_sources(
-                type_stubs,
-                to_remove,
-                base_module,
-                srcs,
-                to_build,
-                to_insert,
-                fs,
-            )
-
-            # python_binary might be the root target, which doesn't have `srcs`.
-            attr_srcs = attrs.get("srcs")
-            attr_srcs = {} if attr_srcs == None else attr_srcs.value()
-            _get_and_materialize_sources(
-                attr_srcs,
-                to_remove,
-                base_module,
-                srcs,
-                to_build,
-                to_insert,
-                fs,
-            )
-
-    if to_build:
-        built = ctx.build(to_build)
-        ensured = ctx.output.ensure_multiple(built)
-
-        # handle individual generated files
-        for srcs, module_name, target in to_insert:
-            _append_or_default_if_module(srcs, module_name, ensured[target][0].rel_path())
-
-        # handle genrules
-        for target, entry in genrule_to_insert:
-            label = target.label.with_sub_target()
-            relative_to = ensured[label][0].rel_path()
-            entry["relative_to"] = relative_to
+            _handle_python_library(entry, attrs, fs, x, known_targets, to_remove, to_build, to_insert)
+
+    _build_generated_files(ctx, to_build, to_insert, genrule_to_insert)
 
     # remove any targets we decided we shouldn't output
     for remove in to_remove: