<commit_before># Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
# Lint as: python3
"""A collection of objects to analyze and visualize Machine Learning Metadata."""

from typing import Any, Dict, ItemsView, KeysView, Optional, Type, Union, ValuesView

from nitroml.analytics import materialized_artifact
from nitroml.analytics import standard_materialized_artifacts
import pandas as pd
from tfx import types

from ml_metadata.google.services.mlmd_service.proto import mlmd_service_pb2
from ml_metadata.google.tfx import metadata_store
from ml_metadata.proto import metadata_store_pb2


def register_standard_artifacts():
  """Registers all artifact classes in 'standard_materialized_artifacts.py'."""
  standard_materialized_artifacts.register_standard_artifacts()


def register_artifact_class(artifact_class:
                            Type[materialized_artifact.MaterializedArtifact]):
  """Registers a subclass of MaterializedArtifact to artifact registery.

  If an artifact_class is registered with the same ARTIFACT_TYPE field,
  it is replaced.

  Args:
    artifact_class: A subclass of MaterializedArtifact to be registered.
  """
  if not issubclass(artifact_class, materialized_artifact.MaterializedArtifact):
    raise ValueError("Artifact class does not subclass "
                     "'materialized_artifact.MaterializedArtifact'")
  materialized_artifact.get_registry().register(artifact_class)


class PropertyDictWrapper:
  """Wraps a dictionary object for property access functionality.

  This class is read-only (setting properties is not implemented).
  """

  def __init__(self, data: Dict[str, Any]):
    self._data = data.copy()

  def __getitem__(self, key):
    return self._data[key]

  def __getattr__(self, key):
    try:
      return self._data[key]
    except KeyError:
      raise AttributeError

  def __contains__(self, val):
    return val in self._data

  def __repr__(self):
    return repr(self._data)

  def get_all(self) -> Dict[str, Any]:
    """Returns dictionary representation of this object."""
    return self._data.copy()

  def keys(self) -> KeysView[str]:
    """Returns list of keys in this object."""
    return self._data.keys()

  def values(self) -> ValuesView[Any]:
    """Returns list of values in this object."""
    return self._data.values()

  def items(self) -> ItemsView[str, Any]:
    """Returns list of key, value pairs in this object."""
    return self._data.items()


class ComponentRun:
  """An object representation of MLMD Component.

  Attributes:
    run_id: The id of the pipeline run that created this component.
  """

  def __init__(self, run_id: str, execution: metadata_store_pb2.Execution,
               store: metadata_store.MetadataStore):
    """Initializes instance of ComponentRun in a given pipeline run.


    Args:
      run_id: The id of the pipeline run that created this component.
      execution: Execution object containing component execution properties.
      store: A store for the artifact metadata.
    """
    self.run_id = run_id
    self._execution = execution
    self._store = store

  @property
  def component_name(self):
    """The name of this component."""
    return self._execution.property['component_id'].string_value

  @property
  def inputs(self) -> PropertyDictWrapper:
    """List of input artifacts in this components workflow.

    Raises:
      NotImplementedError: If this functionality has yet to be implemented.
    """
    raise NotImplementedError(
        'ComponentRun.inputs() is currently unimplemented.')

  @property
  def outputs(self) -> PropertyDictWrapper:
    """List of output artifacts in this components workflow."""
    [component_run_context] = [
        context
        for context in self._store.get_contexts_by_execution(self._execution.id)
        if 'component_id' in context.properties
    ]
    artifacts = self._store.get_artifacts_by_context(component_run_context.id)
    output = {}
    artifact_types = self._store.get_artifact_types_by_id(
        [artifact.type_id for artifact in artifacts])
    for artifact, artifact_type in zip(artifacts, artifact_types):
      tfx_artifact = types.Artifact(artifact_type)
      tfx_artifact.set_mlmd_artifact(artifact)
      materialized_artifact_class = materialized_artifact.get_registry(
      ).get_artifact_class(tfx_artifact.type_name)
      output[tfx_artifact.name] = materialized_artifact_class(tfx_artifact)

    return PropertyDictWrapper(output)

  @property
  def exec_properties(self) -> Dict[str, Any]:
    """A dictionary of user defined exec properties of this component."""
    return self._execution.properties

  @property
  def describe(self) -> pd.Series:
    """Table of system defined execution detail in this component.

    Raises:
      NotImplementedError: If this functionality has yet to be implemented.
    """
    raise NotImplementedError(
        'ComponentRun.describe() is currently unimplemented.')


class PipelineRun:
  """A class representing a single pipeline run.

  Attributes:
    name: The human-readable name of the pipeline.
    run_id: The unique id of the pipeline.
  """

  def __init__(self, name: str, run_id: str, context_id: int,
               store: metadata_store.MetadataStore):
    """Initializes an instance of PipelineRun with an existing run id.

    Args:
      name: The human-readable name of the pipeline.
      run_id: The unique id of the pipeline.<commit_msg>Update component exec properties to unpack value of respective datatype.<commit_after># Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
# Lint as: python3
"""A collection of objects to analyze and visualize Machine Learning Metadata."""

from typing import Any, Dict, ItemsView, KeysView, Optional, Type, Union, ValuesView

from nitroml.analytics import materialized_artifact
from nitroml.analytics import standard_materialized_artifacts
import pandas as pd
from tfx import types

from ml_metadata.google.services.mlmd_service.proto import mlmd_service_pb2
from ml_metadata.google.tfx import metadata_store
from ml_metadata.proto import metadata_store_pb2


def register_standard_artifacts():
  """Registers all artifact classes in 'standard_materialized_artifacts.py'."""
  standard_materialized_artifacts.register_standard_artifacts()


def register_artifact_class(artifact_class:
                            Type[materialized_artifact.MaterializedArtifact]):
  """Registers a subclass of MaterializedArtifact to artifact registery.

  If an artifact_class is registered with the same ARTIFACT_TYPE field,
  it is replaced.

  Args:
    artifact_class: A subclass of MaterializedArtifact to be registered.
  """
  if not issubclass(artifact_class, materialized_artifact.MaterializedArtifact):
    raise ValueError("Artifact class does not subclass "
                     "'materialized_artifact.MaterializedArtifact'")
  materialized_artifact.get_registry().register(artifact_class)


class PropertyDictWrapper:
  """Wraps a dictionary object for property access functionality.

  This class is read-only (setting properties is not implemented).
  """

  def __init__(self, data: Dict[str, Any]):
    self._data = data.copy()

  def __getitem__(self, key):
    return self._data[key]

  def __getattr__(self, key):
    try:
      return self._data[key]
    except KeyError:
      raise AttributeError

  def __contains__(self, val):
    return val in self._data

  def __repr__(self):
    return repr(self._data)

  def get_all(self) -> Dict[str, Any]:
    """Returns dictionary representation of this object."""
    return self._data.copy()

  def keys(self) -> KeysView[str]:
    """Returns list of keys in this object."""
    return self._data.keys()

  def values(self) -> ValuesView[Any]:
    """Returns list of values in this object."""
    return self._data.values()

  def items(self) -> ItemsView[str, Any]:
    """Returns list of key, value pairs in this object."""
    return self._data.items()


class ComponentRun:
  """An object representation of MLMD Component.

  Attributes:
    run_id: The id of the pipeline run that created this component.
  """

  def __init__(self, run_id: str, execution: metadata_store_pb2.Execution,
               store: metadata_store.MetadataStore):
    """Initializes instance of ComponentRun in a given pipeline run.


    Args:
      run_id: The id of the pipeline run that created this component.
      execution: Execution object containing component execution properties.
      store: A store for the artifact metadata.
    """
    self.run_id = run_id
    self._execution = execution
    self._store = store

  @property
  def component_name(self):
    """The name of this component."""
    return self.exec_properties['component_id']

  @property
  def inputs(self) -> PropertyDictWrapper:
    """List of input artifacts in this components workflow.

    Raises:
      NotImplementedError: If this functionality has yet to be implemented.
    """
    raise NotImplementedError(
        'ComponentRun.inputs() is currently unimplemented.')

  @property
  def outputs(self) -> PropertyDictWrapper:
    """List of output artifacts in this components workflow."""
    [component_run_context] = [
        context
        for context in self._store.get_contexts_by_execution(self._execution.id)
        if 'component_id' in context.properties
    ]
    artifacts = self._store.get_artifacts_by_context(component_run_context.id)
    output = {}
    artifact_types = self._store.get_artifact_types_by_id(
        [artifact.type_id for artifact in artifacts])
    for artifact, artifact_type in zip(artifacts, artifact_types):
      tfx_artifact = types.Artifact(artifact_type)
      tfx_artifact.set_mlmd_artifact(artifact)
      materialized_artifact_class = materialized_artifact.get_registry(
      ).get_artifact_class(tfx_artifact.type_name)
      output[tfx_artifact.name] = materialized_artifact_class(tfx_artifact)

    return PropertyDictWrapper(output)

  @property
  def exec_properties(self) -> Dict[str, Any]:
    """A dictionary of user defined exec properties of this component."""
    exec_dict = {}
    for key, val in self._execution.properties.items():
      exec_dict[key] = getattr(val, val.WhichOneof('value'))

    return exec_dict

  @property
  def describe(self) -> pd.Series:
    """Table of system defined execution detail in this component.

    Raises:
      NotImplementedError: If this functionality has yet to be implemented.
    """
    raise NotImplementedError(
        'ComponentRun.describe() is currently unimplemented.')


class PipelineRun:
  """A class representing a single pipeline run.

  Attributes:
    name: The human-readable name of the pipeline.
    run_id: The unique id of the pipeline.
  """

  def __init__(self, name: str, run_id: str, context_id: int,
               store: metadata_store.MetadataStore):
    """Initializes an instance of PipelineRun with an existing run id.

    Args:
      name: The human-readable name of the pipeline.
      run_id: The unique id of the pipeline.
{'fraction_non_alphanumeric': 0.0930955497382199, 'fraction_numerical': 0.001963350785340314, 'mean_word_length': 3.87829209896249, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 19, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5268289', 'n_tokens_mistral': 3430, 'n_tokens_neox': 3218, 'n_words': 1123}