python_arango_ogm.db.pao_edges

 1from __future__ import annotations
 2
 3import os
 4import sys
 5from typing import Dict, Any, Union, Sequence, Type
 6import python_arango_ogm.db
 7from python_arango_ogm.db.pao_db_base import PAODBBase
 8
 9
10class PAOEdgeDef:
11    def __init__(self, from_model: [Type[python_arango_ogm.db.pao_model.PAOModel], str],
12                 to_model: [Type[python_arango_ogm.db.pao_model.PAOModel], str]):
13        self.from_model = from_model
14        self.to_model = to_model
15        if self.to_model is None:
16            raise ValueError("PAOEdgeDef to_model cannot be None")
17
18        app_package = os.getenv('PAO_APP_PACKAGE')
19        if not app_package:
20            raise RuntimeError("PAO_APP_PACKAGE must be defined in the environment (or a .env.test file)")
21        self.models_module_name = f"{app_package}.models"
22
23    def db(self) -> PAODBBase:
24        return self.model_class_from().db
25
26    def insert_edge(
27            self,
28            _from: [str, python_arango_ogm.db.pao_model.PAOModel],
29            _to: Union[str, python_arango_ogm.db.pao_model.PAOModel]
30    ):
31        pass
32
33    def associated_edges(self, lookup_key_dict: Dict[str, Any] = None) -> Sequence[PAOEdgeDef]:
34        from_model_cls = self.model_class_from()
35        to_model_cls = self.model_class_to()
36        from_collection_name = from_model_cls.collection_name()
37        to_collection_name = to_model_cls.collection_name()
38        return from_model_cls.db.get_edge_associations(from_collection_name, to_collection_name, lookup_key_dict)
39
40    def associated_vertices(self, lookup_key_dict: Dict[str, Any], marshall=True) -> Sequence[
41        Type[python_arango_ogm.db.pao_model.PAOModel]]:
42        from_model_cls = self.model_class_from()
43        to_model_cls = self.model_class_to()
44        from_collection_name = from_model_cls.collection_name()
45        to_collection_name = to_model_cls.collection_name()
46        vertices = from_model_cls.db.get_related_vertices(from_collection_name, to_collection_name, lookup_key_dict)
47        return to_model_cls.marshall_row(vertices, marshall)
48
49    def model_class_from(self) -> Type[python_arango_ogm.db.pao_model.PAOModel]:
50        return self.get_model_class(self.from_model)
51
52    def model_class_to(self) -> Type[python_arango_ogm.db.pao_model.PAOModel]:
53        return self.get_model_class(self.to_model)
54
55    @staticmethod
56    def get_model_class(
57            model: Union[str, Type[python_arango_ogm.db.pao_model.PAOModel]]
58    ) -> Type[python_arango_ogm.db.pao_model.PAOModel]:
59        # Get model class, whether it is defined as an actual class or a model:
60        if isinstance(model, str):
61            module = sys.modules[self.models_module_name]
62            result = getattr(module, model)
63        else:
64            result = model
65        return result