WIP: [types] removed Union[str, 'RelationSchema'] to only RelationSchema

yams/schema.py

@@ -255,8 +255,8 @@ class EntitySchema(PermissionMixIn, ERSchema):
 
         if relation_definition is not None:
             # quick access to bounded relation schemas
-            self._subject_relations: Dict[Union[str, "RelationSchema"], "RelationSchema"] = {}
-            self._object_relations: Dict[Union[str, "RelationSchema"], "RelationSchema"] = {}
+            self._subject_relations: Dict["RelationSchema", "RelationSchema"] = {}
+            self._object_relations: Dict["RelationSchema", "RelationSchema"] = {}
             self._specialized_type = relation_definition.specialized_type
             self._specialized_by = relation_definition.specialized_by
             self.final: bool = self.type in BASE_TYPES
@@ -457,18 +457,18 @@ class EntitySchema(PermissionMixIn, ERSchema):
         for relation_schema in self.object_relations():
             yield relation_schema, relation_schema.subjects(self), "object"
 
-    @deprecation.argument_renamed(old_name="rtype", new_name="relation_type")
-    def destination(self, relation_type: Union[str, "RelationSchema"]) -> "EntitySchema":
+    @deprecation.argument_renamed(old_name="rtype", new_name="relation_schema")
+    def destination(self, relation_schema: "RelationSchema") -> "EntitySchema":
         """return the type or schema of entities related by the given subject relation
 
-        `relation_type` is expected to be a non ambiguous relation
+        `relation_schema` is expected to be a non ambiguous relation
         """
-        relation_schema = self._subject_relations[relation_type]
+        relation_schema = self._subject_relations[relation_schema]
         object_types = relation_schema.objects(self.type)
 
         assert len(object_types) == 1, (
             self.type,
-            str(relation_type),
+            str(relation_schema),
             [str(ot) for ot in object_types],
         )
 
@@ -571,7 +571,9 @@ class EntitySchema(PermissionMixIn, ERSchema):
         """return metadata's relation schema if this entity has the given
         `metadata` field for the given `attribute` attribute
         """
-        return self._subject_relations.get("%s_%s" % (attribute, metadata))
+        return self._subject_relations.get(
+            self.schema.relation_schema_for("%s_%s" % (attribute, metadata))
+        )
 
     @deprecation.argument_renamed(old_name="attr", new_name="attribute")
     def is_metadata(self, attribute) -> Optional[Tuple[str, str]]:
@@ -627,9 +629,7 @@ class EntitySchema(PermissionMixIn, ERSchema):
 
     @deprecation.argument_renamed(old_name="skiprels", new_name="skip_relations")
     def is_subobject(
-        self,
-        strict: bool = False,
-        skip_relations: Sequence[Tuple[Union["RelationSchema", str], str]] = (),
+        self, strict: bool = False, skip_relations: Sequence[Tuple["RelationSchema", str]] = (),
     ) -> bool:
         """return True if this entity type is contained by another. If strict,
         return True if this entity type *must* be contained by another.
@@ -673,7 +673,7 @@ class EntitySchema(PermissionMixIn, ERSchema):
 
     def check(
         self,
-        entity: Dict[Union[str, "RelationSchema"], Any],
+        entity: Dict["RelationSchema", Any],
         creation: bool = False,
         _=None,
         relations: Optional[List["RelationSchema"]] = None,
@@ -1020,8 +1020,8 @@ class RelationSchema(ERSchema):
                 self._init_relation(relation_definition)
 
             # mapping to subject/object with schema as key
-            self._subject_schemas: Dict[Any, List] = {}
-            self._object_schemas: Dict[Any, List] = {}
+            self._subject_schemas: Dict["EntitySchema", List["EntitySchema"]] = {}
+            self._object_schemas: Dict["EntitySchema", List["EntitySchema"]] = {}
 
             # relation properties
             self.relation_definitions: Dict[
@@ -1287,9 +1287,7 @@ class RelationSchema(ERSchema):
         return list(self._subject_schemas.items())
 
     @deprecation.argument_renamed(old_name="etype", new_name="entity_type")
-    def subjects(
-        self, entity_type: Optional[Union["EntitySchema", str]] = None
-    ) -> Tuple["EntitySchema", ...]:
+    def subjects(self, entity_type: Optional["EntitySchema"] = None) -> Tuple["EntitySchema", ...]:
         """Return a list of entity schemas which can be subject of this relation.
 
         If entity_type is not None, return a list of schemas which can be subject of
@@ -1306,9 +1304,7 @@ class RelationSchema(ERSchema):
             raise KeyError(f"{self} does not have {entity_type} as object")
 
     @deprecation.argument_renamed(old_name="etype", new_name="entity_type")
-    def objects(
-        self, entity_type: Optional[Union["EntitySchema", str]] = None
-    ) -> Tuple["EntitySchema", ...]:
+    def objects(self, entity_type: Optional["EntitySchema"] = None) -> Tuple["EntitySchema", ...]:
         """Return a list of entity schema which can be object of this relation.
 
         If entity_type is not None, return a list of schemas which can be object of