Merge pull request #227 from linkml/issue-659-pydantic-loader-dumper

Support Pydantic BaseModel in loaders and dumpers

Author: cmungall

linkml_runtime/dumpers/csv_dumper.py

@@ -1,7 +1,8 @@
 import io
 import yaml
 import json
-from typing import Dict, List, Any
+from typing import Union
+from pydantic import BaseModel
 
 from linkml_runtime.dumpers.dumper_root import Dumper
 from linkml_runtime.dumpers.json_dumper import JSONDumper
@@ -15,7 +16,7 @@
 
 class CSVDumper(Dumper):
 
-    def dumps(self, element: YAMLRoot,
+    def dumps(self, element: Union[BaseModel, YAMLRoot],
               index_slot: SlotDefinitionName = None,
               schema: SchemaDefinition = None,
               schemaview: SchemaView = None,

linkml_runtime/dumpers/dumper_root.py

@@ -1,12 +1,14 @@
 from abc import ABC, abstractmethod
+from typing import Union
 
 from linkml_runtime.utils.yamlutils import YAMLRoot
+from pydantic import BaseModel
 
 
 class Dumper(ABC):
     """ Abstract base class for all dumpers """
 
-    def dump(self, element: YAMLRoot, to_file: str, **_) -> None:
+    def dump(self, element: Union[BaseModel, YAMLRoot], to_file: str, **_) -> None:
         """
         Write element to to_file
         :param element: LinkML object to be dumped
@@ -17,10 +19,10 @@ def dump(self, element: YAMLRoot, to_file: str, **_) -> None:
             output_file.write(self.dumps(element, **_))
 
     @abstractmethod
-    def dumps(self, element: YAMLRoot, **_) -> str:
+    def dumps(self, element: Union[BaseModel, YAMLRoot], **_) -> str:
         """
         Convert element to a string
-        @param element: YAMLRoot object to be rendered
+        @param element: Union[BaseModel, YAMLRoot] object to be rendered
         @param _: method specific arguments
         @return: stringified representation of element
         """

linkml_runtime/dumpers/json_dumper.py

@@ -1,6 +1,7 @@
 import json
 from decimal import Decimal
-from typing import Dict
+from typing import Dict, Union
+from pydantic import BaseModel
 
 from deprecated.classic import deprecated
 
@@ -11,9 +12,11 @@
 from linkml_runtime.utils.yamlutils import YAMLRoot, as_json_object
 from jsonasobj2 import JsonObj
 
+
 class JSONDumper(Dumper):
 
-    def dump(self, element: YAMLRoot, to_file: str, contexts: CONTEXTS_PARAM_TYPE = None, **kwargs) -> None:
+    def dump(self, element: Union[BaseModel, YAMLRoot], to_file: str, contexts: CONTEXTS_PARAM_TYPE = None,
+             **kwargs) -> None:
         """
         Write element as json to to_file
         :param element: LinkML object to be serialized as YAML
@@ -26,9 +29,11 @@ def dump(self, element: YAMLRoot, to_file: str, contexts: CONTEXTS_PARAM_TYPE =
             * JSON Object
             * A list containing elements of any type named above
         """
+        if isinstance(element, BaseModel):
+            element = element.dict()
         super().dump(element, to_file, contexts=contexts, **kwargs)
 
-    def dumps(self, element: YAMLRoot, contexts: CONTEXTS_PARAM_TYPE = None, inject_type=True) -> str:
+    def dumps(self, element: Union[BaseModel, YAMLRoot], contexts: CONTEXTS_PARAM_TYPE = None, inject_type=True) -> str:
         """
         Return element as a JSON or a JSON-LD string
         :param element: LinkML object to be emitted
@@ -42,20 +47,24 @@ def dumps(self, element: YAMLRoot, contexts: CONTEXTS_PARAM_TYPE = None, inject_
         :param inject_type: if True (default), add a @type at the top level
         :return: JSON Object representing the element
         """
+
         def default(o):
+            if isinstance(o, BaseModel):
+                return remove_empty_items(o.dict(), hide_protected_keys=True)
             if isinstance(o, YAMLRoot):
                 return remove_empty_items(o, hide_protected_keys=True)
             elif isinstance(o, Decimal):
                 # https://stackoverflow.com/questions/1960516/python-json-serialize-a-decimal-object
                 return str(o)
             else:
                 return json.JSONDecoder().decode(o)
+        if isinstance(element, BaseModel):
+            element = element.dict()
         return json.dumps(as_json_object(element, contexts, inject_type=inject_type),
                           default=default,
                           ensure_ascii=False,
                           indent='  ')
 
-
     @staticmethod
     @deprecated("Use `utils/formatutils/remove_empty_items` instead")
     def remove_empty_items(obj: Dict) -> Dict:
@@ -66,7 +75,8 @@ def remove_empty_items(obj: Dict) -> Dict:
         """
         return formatutils.remove_empty_items(obj, hide_protected_keys=True)
 
-    def to_json_object(self, element: YAMLRoot, contexts: CONTEXTS_PARAM_TYPE = None, inject_type=True) -> JsonObj:
+    def to_json_object(self, element: Union[BaseModel, YAMLRoot], contexts: CONTEXTS_PARAM_TYPE = None,
+                       inject_type=True) -> JsonObj:
         """
         As dumps(), except returns a JsonObj, not a string
 
@@ -83,7 +93,7 @@ def to_json_object(self, element: YAMLRoot, contexts: CONTEXTS_PARAM_TYPE = None
         """
         return as_json_object(element, contexts, inject_type=inject_type)
 
-    def to_dict(self, element: YAMLRoot, **kwargs) -> JsonObj:
+    def to_dict(self, element: Union[BaseModel, YAMLRoot], **kwargs) -> JsonObj:
         """
         As dumps(), except returns a JsonObj, not a string
 
@@ -98,4 +108,4 @@ def to_dict(self, element: YAMLRoot, **kwargs) -> JsonObj:
         :param inject_type: if True (default), add a @type at the top level
         :return: JSON Object representing the element
         """
-        return json.loads(self.dumps(element, inject_type=False, **kwargs))
+        return json.loads(self.dumps(element, inject_type=False, **kwargs))

linkml_runtime/dumpers/rdf_dumper.py

@@ -1,5 +1,6 @@
 import json
-from typing import Optional
+from typing import Optional, Union
+from pydantic import BaseModel
 
 from hbreader import hbread
 from rdflib import Graph
@@ -12,7 +13,7 @@
 
 
 class RDFDumper(Dumper):
-    def as_rdf_graph(self, element: YAMLRoot, contexts: CONTEXTS_PARAM_TYPE, namespaces: CONTEXT_TYPE = None) -> Graph:
+    def as_rdf_graph(self, element: Union[BaseModel, YAMLRoot], contexts: CONTEXTS_PARAM_TYPE, namespaces: CONTEXT_TYPE = None) -> Graph:
         """
         Convert element into an RDF graph guided by the context(s) in contexts
         :param element: element to represent in RDF
@@ -63,7 +64,7 @@ def as_rdf_graph(self, element: YAMLRoot, contexts: CONTEXTS_PARAM_TYPE, namespa
 
         return g
 
-    def dump(self, element: YAMLRoot, to_file: str, contexts: CONTEXTS_PARAM_TYPE = None, fmt: str = 'turtle') -> None:
+    def dump(self, element: Union[BaseModel, YAMLRoot], to_file: str, contexts: CONTEXTS_PARAM_TYPE = None, fmt: str = 'turtle') -> None:
         """
         Write element as rdf to to_file
         :param element: LinkML object to be emitted
@@ -77,15 +78,19 @@ def dump(self, element: YAMLRoot, to_file: str, contexts: CONTEXTS_PARAM_TYPE =
             * A list containing elements of any type named above
         :param fmt: RDF format
         """
+        if isinstance(element, BaseModel):
+            element = element.dict()
         super().dump(element, to_file, contexts=contexts, fmt=fmt)
 
-    def dumps(self, element: YAMLRoot, contexts: CONTEXTS_PARAM_TYPE = None, fmt: Optional[str] = 'turtle') -> str:
+    def dumps(self, element: Union[BaseModel, YAMLRoot], contexts: CONTEXTS_PARAM_TYPE = None, fmt: Optional[str] = 'turtle') -> str:
         """
         Convert element into an RDF graph guided by the context(s) in contexts
         :param element: element to represent in RDF
         :param contexts: JSON-LD context(s) in the form of a file or URL, a json string or a json obj
         :param fmt: rdf format
         :return: rdflib Graph containing element
         """
+        if isinstance(element, BaseModel):
+            element = element.dict()
         return self.as_rdf_graph(remove_empty_items(element, hide_protected_keys=True), contexts).\
             serialize(format=fmt)

linkml_runtime/dumpers/rdflib_dumper.py

@@ -1,7 +1,8 @@
 import logging
 import urllib
 from abc import abstractmethod
-from typing import Optional, Any, Dict
+from typing import Optional, Any, Dict, Union
+from pydantic import BaseModel
 
 from rdflib import Graph, URIRef, XSD
 from rdflib.term import Node, BNode, Literal
@@ -23,7 +24,7 @@ class RDFLibDumper(Dumper):
     This requires a SchemaView object
 
     """
-    def as_rdf_graph(self, element: YAMLRoot, schemaview: SchemaView, prefix_map: Dict[str, str] = None) -> Graph:
+    def as_rdf_graph(self, element: Union[BaseModel, YAMLRoot], schemaview: SchemaView, prefix_map: Dict[str, str] = None) -> Graph:
         """
         Dumps from element to an rdflib Graph,
         following a schema
@@ -136,7 +137,7 @@ def inject_triples(self, element: Any, schemaview: SchemaView, graph: Graph, tar
             graph.add((element_uri, RDF.type, URIRef(schemaview.get_uri(cn, expand=True))))
         return element_uri
 
-    def dump(self, element: YAMLRoot,
+    def dump(self, element: Union[BaseModel, YAMLRoot],
              to_file: str,
              schemaview: SchemaView = None,
              fmt: str = 'turtle', prefix_map: Dict[str, str] = None, **args) -> None:
@@ -152,7 +153,7 @@ def dump(self, element: YAMLRoot,
         """
         super().dump(element, to_file, schemaview=schemaview, fmt=fmt, prefix_map=prefix_map)
 
-    def dumps(self, element: YAMLRoot, schemaview: SchemaView = None,
+    def dumps(self, element: Union[BaseModel, YAMLRoot], schemaview: SchemaView = None,
               fmt: Optional[str] = 'turtle', prefix_map: Dict[str, str] = None) -> str:
         """
         Convert element into an RDF graph guided by the schema

linkml_runtime/dumpers/yaml_dumper.py

@@ -1,18 +1,20 @@
-from decimal import Decimal
-
 import yaml
+from typing import Union
+from pydantic import BaseModel
 
 from linkml_runtime.dumpers.dumper_root import Dumper
 from linkml_runtime.utils.formatutils import remove_empty_items
 from linkml_runtime.utils.yamlutils import YAMLRoot
 
 class YAMLDumper(Dumper):
 
-    def dumps(self, element: YAMLRoot, **kwargs) -> str:
+    def dumps(self, element: Union[BaseModel, YAMLRoot], **kwargs) -> str:
         """ Return element formatted as a YAML string """
         # Internal note: remove_empty_items will also convert Decimals to int/float;
         # this is necessary until https://github.com/yaml/pyyaml/pull/372 is merged
-        return yaml.dump(remove_empty_items(element, hide_protected_keys=True),
+
+        dumper_safe_element = element.dict() if isinstance(element, BaseModel) else element
+        return yaml.dump(remove_empty_items(dumper_safe_element, hide_protected_keys=True),
                          Dumper=yaml.SafeDumper, sort_keys=False, 
                          allow_unicode=True,
                          **kwargs)

linkml_runtime/loaders/csv_loader.py

@@ -2,6 +2,7 @@
 import json
 from typing import Type, Union, List
 from linkml_runtime.utils.yamlutils import YAMLRoot
+from pydantic import BaseModel
 
 from linkml_runtime.loaders.loader_root import Loader
 from linkml_runtime.loaders.json_loader import JSONLoader
@@ -17,7 +18,7 @@ def load_any(self, *args, **kwargs) -> Union[YAMLRoot, List[YAMLRoot]]:
 
 
     def loads(self, input,
-              target_class: Type[YAMLRoot],
+              target_class: Type[Union[BaseModel, YAMLRoot]],
               index_slot: SlotDefinitionName = None,
               schema: SchemaDefinition = None,
               schemaview: SchemaView = None,
@@ -30,7 +31,7 @@ def loads(self, input,
         return JSONLoader().loads(json.dumps({index_slot: objs}), target_class=target_class)
 
     def load(self, source: str,
-             target_class: Type[YAMLRoot],
+             target_class: Type[Union[BaseModel, YAMLRoot]],
              index_slot: SlotDefinitionName = None,
              schema: SchemaDefinition = None,
              schemaview: SchemaView = None,

linkml_runtime/loaders/json_loader.py

@@ -6,12 +6,12 @@
 
 from linkml_runtime.loaders.loader_root import Loader
 from linkml_runtime.utils.yamlutils import YAMLRoot
-
+from pydantic import BaseModel
 
 class JSONLoader(Loader):
 
-    def load_any(self, source: Union[str, dict, TextIO], target_class: Type[YAMLRoot], *, base_dir: Optional[str] = None,
-             metadata: Optional[FileInfo] = None, **_) -> Union[YAMLRoot, List[YAMLRoot]]:
+    def load_any(self, source: Union[str, dict, TextIO], target_class: Type[Union[BaseModel, YAMLRoot]], *, base_dir: Optional[str] = None,
+             metadata: Optional[FileInfo] = None, **_) -> Union[BaseModel, YAMLRoot, List[BaseModel], List[YAMLRoot]]:
         def loader(data: Union[str, dict], _: FileInfo) -> Optional[Dict]:
             data_as_dict = json.loads(data) if isinstance(data, str) else data
             if isinstance(data_as_dict, list):

linkml_runtime/loaders/loader_root.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import TextIO, Union, Optional, Callable, Dict, Type, Any, List
 
+from pydantic import BaseModel
 from hbreader import FileInfo, hbread
 from jsonasobj2 import as_dict, JsonObj
 
@@ -36,9 +37,9 @@ def _is_empty(o) -> bool:
     def load_source(self,
                     source: Union[str, dict, TextIO],
                     loader: Callable[[Union[str, Dict], FileInfo], Optional[Union[Dict, List]]],
-                    target_class: Type[YAMLRoot],
+                    target_class: Union[Type[YAMLRoot], Type[BaseModel]],
                     accept_header: Optional[str] = "text/plain, application/yaml;q=0.9",
-                    metadata: Optional[FileInfo] = None) -> Optional[Union[YAMLRoot, List[YAMLRoot]]]:
+                    metadata: Optional[FileInfo] = None) -> Optional[Union[BaseModel, YAMLRoot, List[BaseModel], List[YAMLRoot]]]:
         """ Base loader - convert a file, url, string, open file handle or dictionary into an instance
         of target_class
 
@@ -59,19 +60,28 @@ def load_source(self,
         else:
             data = source
         data_as_dict = loader(data, metadata)
+
         if data_as_dict:
             if isinstance(data_as_dict, list):
-                return [target_class(**as_dict(x)) for x in data_as_dict]
+               if issubclass(target_class, YAMLRoot):
+                   return [target_class(**as_dict(x)) for x in data_as_dict]
+               elif issubclass(target_class, BaseModel):
+                   return [target_class.parse_obj(**as_dict(x)) for x in data_as_dict]
+               else:
+                   raise ValueError(f'Cannot load list of {target_class}')
             elif isinstance(data_as_dict, dict):
-                return target_class(**data_as_dict)
+                if issubclass(target_class, BaseModel):
+                    return target_class.parse_obj(data_as_dict)
+                else:
+                    return target_class(**data_as_dict)
             elif isinstance(data_as_dict, JsonObj):
                 return [target_class(**as_dict(x)) for x in data_as_dict]
             else:
                 raise ValueError(f'Unexpected type {data_as_dict}')
         else:
             return None
 
-    def load(self, *args, **kwargs) -> YAMLRoot:
+    def load(self, *args, **kwargs) -> Union[BaseModel, YAMLRoot]:
         """
         Load source as an instance of target_class
 
@@ -83,14 +93,14 @@ def load(self, *args, **kwargs) -> YAMLRoot:
         :return: instance of target_class
         """
         results = self.load_any(*args, **kwargs)
-        if isinstance(results, YAMLRoot):
+        if isinstance(results, BaseModel) or isinstance(results, YAMLRoot):
             return results
         else:
-            raise ValueError(f'Result is not an instance of YAMLRoot: {type(results)}')
+            raise ValueError(f'Result is not an instance of BaseModel or YAMLRoot: {type(results)}')
 
     @abstractmethod
-    def load_any(self, source: Union[str, dict, TextIO], target_class: Type[YAMLRoot], *, base_dir: Optional[str] = None,
-             metadata: Optional[FileInfo] = None, **_) -> Union[YAMLRoot, List[YAMLRoot]]:
+    def load_any(self, source: Union[str, dict, TextIO], target_class: Type[Union[BaseModel, YAMLRoot]], *, base_dir: Optional[str] = None,
+             metadata: Optional[FileInfo] = None, **_) -> Union[BaseModel, YAMLRoot, List[BaseModel], List[YAMLRoot]]:
         """
         Load source as an instance of target_class, or list of instances of target_class
 
@@ -103,7 +113,7 @@ def load_any(self, source: Union[str, dict, TextIO], target_class: Type[YAMLRoot
         """
         raise NotImplementedError()
 
-    def loads_any(self, source: str, target_class: Type[YAMLRoot], *, metadata: Optional[FileInfo] = None, **_) -> Union[YAMLRoot, List[YAMLRoot]]:
+    def loads_any(self, source: str, target_class: Type[Union[BaseModel, YAMLRoot]], *, metadata: Optional[FileInfo] = None, **_) -> Union[BaseModel, YAMLRoot, List[BaseModel], List[YAMLRoot]]:
         """
         Load source as a string as an instance of target_class, or list of instances of target_class
         @param source: source
@@ -114,7 +124,7 @@ def loads_any(self, source: str, target_class: Type[YAMLRoot], *, metadata: Opti
         """
         return self.load_any(source, target_class, metadata=metadata)
 
-    def loads(self, source: str, target_class: Type[YAMLRoot], *, metadata: Optional[FileInfo] = None, **_) -> YAMLRoot:
+    def loads(self, source: str, target_class: Type[Union[BaseModel, YAMLRoot]], *, metadata: Optional[FileInfo] = None, **_) -> Union[BaseModel, YAMLRoot]:
         """
         Load source as a string
         :param source: source

linkml_runtime/loaders/rdf_loader.py

@@ -5,6 +5,7 @@
 from linkml_runtime.loaders.loader_root import Loader
 from linkml_runtime.utils.context_utils import CONTEXTS_PARAM_TYPE
 from linkml_runtime.utils.yamlutils import YAMLRoot
+from pydantic import BaseModel
 from rdflib import Graph
 
 from linkml_runtime.loaders.requests_ssl_patch import no_ssl_verification
@@ -15,12 +16,13 @@
 
 class RDFLoader(Loader):
 
-    def load_any(self, *args, **kwargs) -> Union[YAMLRoot, List[YAMLRoot]]:
+    def load_any(self, *args, **kwargs) -> Union[BaseModel, YAMLRoot, List[BaseModel], List[YAMLRoot]]:
         return self.load(*args, **kwargs)
 
-    def load(self, source: Union[str, TextIO, Graph], target_class: Type[YAMLRoot], *, base_dir: Optional[str] = None,
+
+    def load(self, source: Union[str, TextIO, Graph], target_class: Type[Union[BaseModel, YAMLRoot]], *, base_dir: Optional[str] = None,
              contexts: CONTEXTS_PARAM_TYPE = None, fmt: Optional[str] = 'turtle',
-             metadata: Optional[FileInfo] = None) -> YAMLRoot:
+             metadata: Optional[FileInfo] = None) -> Union[BaseModel, YAMLRoot]:
         """
         Load the RDF in source into the python target_class structure
         :param source: RDF data source. Can be a URL, a file name, an RDF string, an open handle or an existing graph