pantsbuild / pants / 19015773527

from __future__ import annotations

import logging

from abc import ABCMeta

from dataclasses import dataclass

from pants.core.environments.rules import EnvironmentNameRequest, resolve_environment_name

from pants.engine.environment import EnvironmentName

from pants.engine.fs import MergeDigests, Snapshot, Workspace

from pants.engine.goal import Goal, GoalSubsystem

from pants.engine.internals.specs_rules import find_valid_field_sets_for_target_roots

from pants.engine.intrinsics import digest_to_snapshot

from pants.engine.rules import collect_rules, concurrently, goal_rule, implicitly, rule

from pants.engine.target import FieldSet, NoApplicableTargetsBehavior, TargetRootsToFieldSetsRequest

from pants.engine.unions import UnionMembership, union

logger = logging.getLogger(__name__)

@union(in_scope_types=[EnvironmentName])

class GenerateSnapshotsFieldSet(FieldSet, metaclass=ABCMeta):

@dataclass(frozen=True)

class GenerateSnapshotsResult:

    snapshot: Snapshot

@rule(polymorphic=True)

async def generate_snapshots(

    raise NotImplementedError()

@dataclass(frozen=True)

class EnvironmentAwareGenerateSnapshotsRequest:

    field_set: GenerateSnapshotsFieldSet

@rule

async def environment_aware_generate_snapshots(

    environment_name = await resolve_environment_name(

    result = await generate_snapshots(

    return result

class GenerateSnapshotsSubsystem(GoalSubsystem):

    name = "generate-snapshots"

    help = "Generate test snapshots."

    @classmethod

    def activated(cls, union_membership: UnionMembership) -> bool:

        return GenerateSnapshotsFieldSet in union_membership

class GenerateSnapshots(Goal):

    subsystem_cls = GenerateSnapshotsSubsystem

    environment_behavior = Goal.EnvironmentBehavior.USES_ENVIRONMENTS

@goal_rule

async def generate_snapshots_goal(workspace: Workspace) -> GenerateSnapshots:

    target_roots_to_field_sets = await find_valid_field_sets_for_target_roots(

    if not target_roots_to_field_sets.field_sets:

        return GenerateSnapshots(exit_code=0)

    snapshot_results = await concurrently(

    all_snapshots = await digest_to_snapshot(

    workspace.write_digest(all_snapshots.digest)

    for file in all_snapshots.files:

        logger.info(f"Generated {file}")

    return GenerateSnapshots(exit_code=0)

def rules():

    return collect_rules()