import pathlib

import importlib

import contextlib

from typing import List, Dict, Any, Union

from ..base import BaseConfig

from ..df.base import BaseOrchestrator, OperationImplementation, opimp_in

from ..df.types import DataFlow, Input, Operation, Stage, FailedToLoadOperation

from ..df.exceptions import DefinitionNotFoundInDataFlow

from ..df.memory import (

MemoryOrchestrator,

MemoryInputSet,

MemoryInputSetConfig,

StringInputSetContext,

)

from ..configloader.configloader import BaseConfigLoader

from ..configloader.json import JSONConfigLoader

from ..source.source import SubsetSources, Sources

from ..util.data import merge, split_dot_seperated, traverse_set

from ..util.entrypoint import load

from ..util.cli.cmd import CMD, CMDOutputOverride, DFFMLCLICMD

from ..util.cli.cmds import (

SourcesCMD,

KeysCMD,

KeysCMDConfig,

)

from ..overlay.overlay import merge_implementations

from ..util.cli.parser import ParseInputsAction

from ..util.config.fields import FIELD_SOURCES

from ..util.crypto import insecure_hash

from ..base import config, field

from ..high_level.dataflow import run as run_dataflow

@config

class MergeConfig:

dataflows: List[pathlib.Path] = field("DataFlows to merge")

configloader: BaseConfigLoader = field(

"ConfigLoader to use for exporting", default=JSONConfigLoader,

)

not_linked: bool = field(

"Do not export dataflows as linked", default=False,

)

class Merge(CMD):

CONFIG = MergeConfig

async def run(self):

# The merged dataflow

merged: Dict[str, Any] = {}

# For entering ConfigLoader contexts

async with contextlib.AsyncExitStack() as exit_stack:

# Load config loaders we'll need as we see their file types

parsers: Dict[str, BaseConfigLoader] = {}

for path in self.dataflows:

_, exported = await BaseConfigLoader.load_file(

parsers, exit_stack, path

)

merge(merged, exported, list_append=True)

# Export the dataflow

dataflow = DataFlow._fromdict(**merged)

async with self.configloader(BaseConfig()) as configloader:

async with configloader() as loader:

exported = dataflow.export(linked=not self.not_linked)

print((await loader.dumpb(exported)).decode())

@config

class CreateConfig:

operations: List[str] = field("Operations to create a dataflow for",)

configloader: BaseConfigLoader = field(

"ConfigLoader to use", default=JSONConfigLoader,

)

not_linked: bool = field(

"Do not export dataflows as linked", default=False,

)

inputs: List[str] = field(

"Inputs to be added to every context",

action=ParseInputsAction,

default_factory=lambda: [],

)

flow: List[str] = field(

"Flow of inputs", action=ParseInputsAction, default_factory=lambda: [],

)

config: List[str] = field(

"configs", action=ParseInputsAction, default_factory=lambda: [],

)

class Create(CMD):

CONFIG = CreateConfig

async def run(self):

operations = []

operation_names = []

for load_operation in self.operations:

if "=" in load_operation:

opname, operation = load_operation.split("=")

else:

opname = operation = load_operation

operation_names.append(opname)

if ":" in operation:

operations.extend(

map(

OperationImplementation._imp,

load(operation, relative=True),

)

)

else:

try:

operations += [Operation.load(operation)]

except FailedToLoadOperation:

operations += opimp_in(importlib.import_module(operation))

async with self.configloader(BaseConfig()) as configloader:

async with configloader() as loader:

dataflow = DataFlow(

operations=dict(zip(operation_names, operations))

)

# flow argument key is of the form opname.inputs/conditions.keyname

for v, k in self.flow:

*opname, val_type, key = split_dot_seperated(k)

opname = ".".join(opname)

if val_type == "inputs":

dataflow.flow[opname].inputs[key] = v

else:

dataflow.flow[opname].conditions = v

for v, k in self.config:

traverse_set(dataflow.configs, k, value=v)

exported = dataflow.export(linked=not self.not_linked)

if self.inputs:

if not "seed" in exported:

exported["seed"] = []

for val, input_info in self.inputs:

if "=" in input_info:

def_name, origin = input_info.split(

"=", maxsplit=1

)

# self.inputs is of the form val=def_name=origin

exported["seed"].append(

{

"value": val,

"definition": def_name,

"origin": origin,

}

)

else:

# self.inputs is of the form val=def_name

exported["seed"].append(

{"value": val, "definition": input_info}

)

print((await loader.dumpb(exported)).decode())

@config

class RunCMDConfig:

dataflow: Union[str, DataFlow] = field(

"File containing exported DataFlow",

)

configloader: BaseConfigLoader = field(

"ConfigLoader to use for importing DataFlow", default=None,

)

sources: Sources = FIELD_SOURCES

caching: List[str] = field(

"Skip running DataFlow if a record already contains these features",

default_factory=lambda: [],

)

passcmd: bool = field(

"Pass the instance DFFML.CLI command to the flow", default=False,

)

no_update: bool = field(

"Update record with sources", default=False,

)

no_echo: bool = field(

"Do not echo back records", default=False,

)

no_strict: bool = field(

"Do not exit on operation exceptions, just log errors", default=False,

)

orchestrator: BaseOrchestrator = field(

"Orchestrator", default=MemoryOrchestrator,

)

inputs: List[str] = field(

"Other inputs to add under each ctx (record's key will "

+ "be used as the context)",

action=ParseInputsAction,

default_factory=lambda: [],

)

record_def: str = field(

"Definition to be used for record.key."

+ "If set, record.key will be added to the set of inputs "

+ "under each context (which is also the record's key)",

default=False,

)

config: List[str] = field(

"configs", action=ParseInputsAction, default_factory=lambda: [],

)

class RunCMD(SourcesCMD):

CONFIG = RunCMDConfig

def __init__(self, *args, **kwargs):

super().__init__(*args, **kwargs)

self.orchestrator = self.orchestrator.withconfig(self.extra_config)

@config

class RunAllRecordsConfig(RunCMDConfig):

pass

class RunAllRecords(RunCMD):

"""Run dataflow for all records in sources"""

CONFIG = RunAllRecordsConfig

async def records(self, sctx):

"""

This method exists so that it can be overriden by RunRecordSet

"""

async for record in sctx.records():

yield record

async def run_dataflow(self, orchestrator, sources, dataflow):

# Orchestrate the running of these operations

async with orchestrator(dataflow) as octx, sources() as sctx:

# Add our inputs to the input network with the context being the

# record key

inputs = []

async for record in self.records(sctx):

# Skip running DataFlow if record already has features

existing_features = record.features()

if self.caching and all(

map(

lambda cached: cached in existing_features,

self.caching,

)

):

continue

record_inputs = []

for value, def_name in self.inputs:

record_inputs.append(

Input(

value=value,

definition=dataflow.definitions[def_name],

)

)

if self.record_def:

record_inputs.append(

Input(

value=record.key,

definition=dataflow.definitions[self.record_def],

)

)

if self.passcmd:

record_inputs.append(

Input(

value=self,

definition=DFFMLCLICMD,

)

)

# TODO(p1) When OrchestratorContext is fixed to accept an

# asyncgenerator we won't have to build a list

inputs.append(

MemoryInputSet(

MemoryInputSetConfig(

ctx=StringInputSetContext(record.key),

inputs=record_inputs,

)

)

)

if not inputs:

return

async for ctx, results in octx.run(

*inputs, strict=not self.no_strict

):

ctx_str = (await ctx.handle()).as_string()

# TODO(p4) Make a RecordInputSetContext which would let us

# store the record instead of recalling it by the URL

record = await sctx.record(ctx_str)

# Store the results

record.evaluated(results)

yield record

if not self.no_update:

await sctx.update(record)

async def run(self):

async with self.orchestrator as orchestrator, self.sources as sources:

async for record in self.run_dataflow(

orchestrator, sources, self.dataflow

):

if not self.no_echo:

yield record

if self.no_echo:

yield CMDOutputOverride

async def __aenter__(self):

if not isinstance(self.dataflow, DataFlow):

dataflow_path = pathlib.Path(self.dataflow)

config_cls = self.configloader

if config_cls is None:

config_type = dataflow_path.suffix.replace(".", "")

config_cls = BaseConfigLoader.load(config_type)

async with config_cls.withconfig(

self.extra_config

) as configloader:

async with configloader() as loader:

exported = await loader.loadb(dataflow_path.read_bytes())

self.dataflow = DataFlow._fromdict(**exported)

for v, k in self.config:

traverse_set(dataflow.configs, k, value=v)

return self

@config

class RunRecordSetConfig(RunAllRecordsConfig, KeysCMDConfig):

pass

class RunRecordSet(RunAllRecords, KeysCMD):

"""Run dataflow for single record or set of records"""

CONFIG = RunRecordSetConfig

async def records(self, sctx):

for key in self.keys:

yield await sctx.record(key)

def __init__(self, *args, **kwargs):

super().__init__(*args, **kwargs)

self.sources = SubsetSources(*self.sources, keys=self.keys)

class RunRecords(CMD):

"""Run DataFlow and assign output to a record"""

_set = RunRecordSet

_all = RunAllRecords

@config

class RunSingleConfig:

dataflow: Union[str, DataFlow] = field(

"File containing exported DataFlow",

)

overlay: Union[str, DataFlow] = field(

"File containing exported DataFlow to use as an overlay", default=None,

)

no_echo: bool = field(

"Do not echo back records", default=False,

)

configloader: BaseConfigLoader = field(

"ConfigLoader to use for importing DataFlow", default=None,

)

orchestrator: BaseOrchestrator = field(

"Orchestrator", default=MemoryOrchestrator,

)

inputs: List[str] = field(

"Other inputs to add under each ctx",

action=ParseInputsAction,

default_factory=lambda: [],

)

no_strict: bool = field(

"Do not exit on operation exceptions, just log errors", default=False,

)

config: List[str] = field(

"configs", action=ParseInputsAction, default_factory=lambda: [],

)

class RunSingle(CMD):

CONFIG = RunSingleConfig

def __init__(self, *args, **kwargs):

super().__init__(*args, **kwargs)

self.orchestrator = self.orchestrator.withconfig(self.extra_config)

async def get_dataflow(self, dataflow_path):

if isinstance(dataflow_path, DataFlow):

return dataflow_path

if not isinstance(dataflow_path, pathlib.Path):

dataflow_path = pathlib.Path(dataflow_path)

config_cls = self.configloader

if config_cls is None:

config_type = dataflow_path.suffix.replace(".", "")

config_cls = BaseConfigLoader.load(config_type)

async with config_cls.withconfig(self.extra_config) as configloader:

async with configloader() as loader:

exported = await loader.loadb(dataflow_path.read_bytes())

dataflow = DataFlow._fromdict(**exported)

for v, k in self.config:

traverse_set(dataflow.configs, k, value=v)

return dataflow

async def get_overlay(self, overlay_path):

if isinstance(overlay_path, DataFlow):

return overlay_path

if not isinstance(overlay_path, (str, pathlib.Path)):

return overlay_path

if isinstance(overlay_path, str):

overlay_path = pathlib.Path(overlay_path)

config_cls = self.configloader

if config_cls is None:

config_type = overlay_path.suffix.replace(".", "")

config_cls = BaseConfigLoader.load(config_type)

async with config_cls.withconfig(self.extra_config) as configloader:

async with configloader() as loader:

exported = await loader.loadb(overlay_path.read_bytes())

overlay = DataFlow._fromdict(**exported)

for v, k in self.config:

traverse_set(overlay.configs, k, value=v)

return overlay

def input_objects(self, dataflow):

for value, def_name in self.inputs:

if not def_name in dataflow.definitions:

raise DefinitionNotFoundInDataFlow(

f"{def_name!r} not found in {list(dataflow.definitions.keys())}"

)

yield Input(value=value, definition=dataflow.definitions[def_name])

async def run(self):

dataflow = await self.get_dataflow(self.dataflow)

overlay = await self.get_overlay(self.overlay)

if overlay is not None:

unoverlayed_dataflow = dataflow

dataflow = DataFlow._fromdict(

**merge(unoverlayed_dataflow.export(), overlay.export(),)

)

dataflow = merge_implementations(

dataflow, unoverlayed_dataflow, overlay,

)

# TODO We have an issue with ``links``

async for ctx, results in run_dataflow(

dataflow,

list(self.input_objects(dataflow)),

orchestrator=self.orchestrator,

strict=not self.no_strict,

):

if not self.no_echo:

yield results

if self.no_echo:

yield CMDOutputOverride

@config

class RunContextsConfig(RunSingleConfig):

context_def: str = field(

"Definition to be used for contexts key. "

+ "If set, the key will be added to the set of inputs "

+ "under each context (which is also the contexts name)",

default=False,

)

contexts: List[str] = field(

"Contexts to run", default_factory=lambda: ["context1"], required=False

)

class RunContexts(RunSingle):

CONFIG = RunContextsConfig

async def run(self):

dataflow = await self.get_dataflow(self.dataflow)

common_inputs = list(self.input_objects(dataflow))

if not self.context_def in dataflow.definitions:

raise DefinitionNotFoundInDataFlow(

f"{self.context_def!r} not found in {list(dataflow.definitions.keys())}"

)

dataflow_inputs = {

ctx_string: [

Input(

value=ctx_string,

definition=dataflow.definitions[self.context_def],

)

]

+ common_inputs

for ctx_string in self.contexts

}

async for ctx, result in run_dataflow(

dataflow,

dataflow_inputs,

orchestrator=self.orchestrator,

strict=not self.no_strict,

):

if not self.no_echo:

yield {(await ctx.handle()).as_string(): result}

if self.no_echo:

yield CMDOutputOverride

class Run(CMD):

"""Run dataflow"""

single = RunSingle

contexts = RunContexts

records = RunRecords

@config

class DiagramConfig:

dataflow: Union[str, DataFlow] = field(

"File containing exported DataFlow",

)

configloader: BaseConfigLoader = field(

"ConfigLoader to use for importing DataFlow", default=None,

)

stages: List[str] = field(

"Which stages to display: (processing, cleanup, output)",

default_factory=lambda: [],

)

simple: bool = field("Don't display input and output names", default=False)

shortname: bool = field("Cut long operation names", default=False)

display: str = field(

"How to display (TD: top down, LR, RL, BT)", default="TD",

)

class Diagram(CMD):

CONFIG = DiagramConfig

async def __aenter__(self):

if not isinstance(self.dataflow, DataFlow):

dataflow_path = pathlib.Path(self.dataflow)

config_cls = self.configloader

if config_cls is None:

config_type = dataflow_path.suffix.replace(".", "")

config_cls = BaseConfigLoader.load(config_type)

async with config_cls.withconfig(

self.extra_config

) as configloader:

async with configloader() as loader:

exported = await loader.loadb(dataflow_path.read_bytes())

self.dataflow = DataFlow._fromdict(**exported)

return self

async def run(self):

dataflow = self.dataflow

print(f"graph {self.display}")

for stage in Stage:

# Skip stage if not wanted

if self.stages and stage.value not in self.stages:

continue

stage_node = insecure_hash("stage." + stage.value)

if len(self.stages) != 1:

print(f"subgraph {stage_node}[{stage.value.title()} Stage]")

print(f"style {stage_node} fill:#afd388b5,stroke:#a4ca7a")

for instance_name, operation in dataflow.operations.items():

if operation.stage != stage:

continue

subgraph_node = insecure_hash("subgraph." + instance_name)

node = insecure_hash(instance_name)

if not self.simple:

if ":" in instance_name and self.shortname:

print(

f"subgraph {subgraph_node}[{instance_name.split(':')[-1]}]"

)

else:

print(f"subgraph {subgraph_node}[{instance_name}]")

print(f"style {subgraph_node} fill:#fff4de,stroke:#cece71")

if ":" in operation.instance_name and self.shortname:

print(f"{node}[{operation.instance_name.split(':')[-1]}]")

else:

print(f"{node}[{operation.instance_name}]")

for input_name in operation.inputs.keys():

input_node = insecure_hash(

"input." + instance_name + "." + input_name

)

if not self.simple:

print(f"{input_node}({input_name})")

print(f"{input_node} --> {node}")

for output_name in operation.outputs.keys():

output_node = insecure_hash(

"output." + instance_name + "." + output_name

)

if not self.simple:

print(f"{output_node}({output_name})")

print(f"{node} --> {output_node}")

for condition in operation.conditions:

condition_node = insecure_hash(

"condition." + instance_name + "." + condition.name

)

if not self.simple:

print(f"{condition_node}{'{' + condition.name + '}'}")

print(f"{condition_node} --> {node}")

if not self.simple:

print(f"end")

if len(self.stages) != 1:

print(f"end")

if len(self.stages) != 1:

print(f"subgraph inputs[Inputs]")

print(f"style inputs fill:#f6dbf9,stroke:#a178ca")

for instance_name, input_flow in dataflow.flow.items():

operation = dataflow.operations[instance_name]

if self.stages and not operation.stage.value in self.stages:

continue

node = insecure_hash(instance_name)

for input_name, sources in input_flow.inputs.items():

for source in sources:

# TODO Put various sources in their own "Inputs" subgraphs

if isinstance(source, str):

input_definition = operation.inputs[input_name]

seed_input_node = insecure_hash(

source + "." + input_definition.name

)

print(

f"{seed_input_node}({source}<br>{input_definition.name})"

)

if len(self.stages) == 1:

print(

f"style {seed_input_node} fill:#f6dbf9,stroke:#a178ca"

)

if not self.simple:

input_node = insecure_hash(

"input." + instance_name + "." + input_name

)

print(f"{seed_input_node} --> {input_node}")

else:

print(f"{seed_input_node} --> {node}")

elif isinstance(source, dict) and isinstance(

list(source.values())[0], list

):

# Handle the case where the key is the origin (so

# something like "seed") and the value is a list of

# acceptable definitions from that origin.

origin, definition_names = list(source.items())[0]

for definition_name in definition_names:

origin_definition_name = (

origin + "." + definition_name

)

seed_input_node = insecure_hash(

origin_definition_name

)

print(

f"{seed_input_node}({origin}<br>{definition_name})"

)

if len(self.stages) == 1:

print(

f"style {seed_input_node} fill:#f6dbf9,stroke:#a178ca"

)

if not self.simple:

input_node = insecure_hash(

"input." + instance_name + "." + input_name

)

print(f"{seed_input_node} --> {input_node}")

else:

print(f"{seed_input_node} --> {node}")

else:

# In order to support selection an input based using an

# alternate definition along with restriction to inputs

# who's origins match the alternate definitions in the

# list. We select the first output source since that

# will be the immediate alternate definition

if (

isinstance(source, list)

and source

and isinstance(source[0], dict)

):

source = source[0]

if not self.simple:

source_output_node = insecure_hash(

"output." + ".".join(list(source.items())[0])

)

input_node = insecure_hash(

"input." + instance_name + "." + input_name

)

print(f"{source_output_node} --> {input_node}")

else:

source_operation_node = insecure_hash(

list(source.keys())[0]

)

print(f"{source_operation_node} --> {node}")

for i, condition in enumerate(input_flow.conditions):

if isinstance(condition, str):

if not self.simple:

condition_name = operation.conditions[i].name

seed_condition_node = insecure_hash(

condition + "." + condition_name

)

print(f"{seed_condition_node}({condition_name})")

seed_dependent_node = insecure_hash(

"condition." + instance_name + "." + condition_name

)

print(

f"{seed_condition_node} --> {seed_dependent_node}"

)

else:

if not self.simple:

dependee_node = insecure_hash(

"output." + ".".join(list(condition.items())[0])

)

dependent_node = insecure_hash(

"condition."

+ instance_name

+ "."

+ dataflow.operations[list(condition.keys())[0]]

.outputs[list(condition.values())[0]]

.name

)

print(f"{dependee_node} --> {dependent_node}")

else:

dependee_operation_node = insecure_hash(

list(condition.keys())[0]

)

print(f"{dependee_operation_node} --> {node}")

if len(self.stages) != 1:

print(f"end")

# Name collision

class Dataflow(CMD):

merge = Merge

create = Create

run = Run

diagram = Diagram