def _get_artifact_type()

in tfx/types/artifact.py [0:0]

• 34 lines of code
• 14 McCabe index (conditional complexity)

  def _get_artifact_type(cls):
    existing_artifact_type = getattr(cls, '_MLMD_ARTIFACT_TYPE', None)
    if (not existing_artifact_type) or (cls.TYPE_NAME !=
                                        existing_artifact_type.name):
      type_name = cls.TYPE_NAME
      if not (type_name and isinstance(type_name, str)):
        raise ValueError(
            ('The Artifact subclass %s must override the TYPE_NAME attribute '
             'with a string type name identifier (got %r instead).') %
            (cls, type_name))
      artifact_type = metadata_store_pb2.ArtifactType()
      artifact_type.name = type_name

      # Populate ML Metadata artifact properties dictionary.
      if cls.PROPERTIES:
        # Perform validation on PROPERTIES dictionary.
        if not isinstance(cls.PROPERTIES, dict):
          raise ValueError(
              'Artifact subclass %s.PROPERTIES is not a dictionary.' % cls)
        for key, value in cls.PROPERTIES.items():
          if not (isinstance(key,
                             (str, bytes)) and isinstance(value, Property)):
            raise ValueError(
                ('Artifact subclass %s.PROPERTIES dictionary must have keys of '
                 'type string and values of type artifact.Property.') % cls)

        for key, value in cls.PROPERTIES.items():
          artifact_type.properties[key] = value.mlmd_type()

      # Populate ML Metadata artifact type field: `base_type`.
      type_annotation_cls = cls.TYPE_ANNOTATION
      if type_annotation_cls:
        if not issubclass(type_annotation_cls, SystemArtifact):
          raise ValueError(
              'TYPE_ANNOTATION %s is not a subclass of SystemArtifact.' %
              type_annotation_cls)
        if type_annotation_cls.MLMD_SYSTEM_BASE_TYPE:
          artifact_type.base_type = type_annotation_cls.MLMD_SYSTEM_BASE_TYPE

      cls._MLMD_ARTIFACT_TYPE = artifact_type
    return copy.deepcopy(cls._MLMD_ARTIFACT_TYPE)