lamindb 0.77.2__py3-none-any.whl → 1.0rc1__py3-none-any.whl

lamindb/core/_data.py

@@ -1,26 +1,26 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from django.db import connections
 from lamin_utils import colors, logger
 from lamindb_setup.core._docs import doc_args
-from lnschema_core.models import (
+
+from lamindb._query_set import QuerySet
+from lamindb.core._settings import settings
+from lamindb.models import (
     Artifact,
     Collection,
     Feature,
-    FeatureSet,
     Record,
     Run,
+    Schema,
     ULabel,
     format_field_value,
     record_repr,
 )
 
-from lamindb._query_set import QuerySet
-from lamindb.core._settings import settings
-
 from ._context import context
 from ._django import get_artifact_with_related, get_related_model
 from ._feature_manager import (
@@ -29,15 +29,15 @@ from ._feature_manager import (
     get_label_links,
 )
 from .exceptions import ValidationError
-from .schema import (
+from .relations import (
+    dict_module_name_to_model_name,
     dict_related_model_to_related_name,
-    dict_schema_name_to_model_name,
 )
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
-    from lnschema_core.types import StrField
+    from lamindb.base.types import StrField
 
 
 WARNING_RUN_TRANSFORM = "no run & transform got linked, call `ln.track()` & re-run"
@@ -56,49 +56,44 @@ def get_run(run: Run | None) -> Run | None:
     return run
 
 
-def add_transform_to_kwargs(kwargs: dict[str, Any], run: Run):
-    if run is not None:
-        kwargs["transform"] = run.transform
-
-
-def save_feature_sets(self: Artifact | Collection) -> None:
-    if hasattr(self, "_feature_sets"):
-        from lamindb.core._feature_manager import get_feature_set_by_slot_
+def save_staged__schemas_m2m(self: Artifact | Collection) -> None:
+    if hasattr(self, "_staged__schemas_m2m"):
+        from lamindb.core._feature_manager import get_schema_by_slot_
 
-        existing_feature_sets = get_feature_set_by_slot_(self)
-        saved_feature_sets = {}
-        for key, feature_set in self._feature_sets.items():
-            if isinstance(feature_set, FeatureSet) and feature_set._state.adding:
-                feature_set.save()
-                saved_feature_sets[key] = feature_set
-            if key in existing_feature_sets:
+        existing_staged__schemas_m2m = get_schema_by_slot_(self)
+        saved_staged__schemas_m2m = {}
+        for key, schema in self._staged__schemas_m2m.items():
+            if isinstance(schema, Schema) and schema._state.adding:
+                schema.save()
+                saved_staged__schemas_m2m[key] = schema
+            if key in existing_staged__schemas_m2m:
                 # remove existing feature set on the same slot
-                self.feature_sets.remove(existing_feature_sets[key])
-        if len(saved_feature_sets) > 0:
-            s = "s" if len(saved_feature_sets) > 1 else ""
-            display_feature_set_keys = ",".join(
-                f"'{key}'" for key in saved_feature_sets.keys()
+                self._schemas_m2m.remove(existing_staged__schemas_m2m[key])
+        if len(saved_staged__schemas_m2m) > 0:
+            s = "s" if len(saved_staged__schemas_m2m) > 1 else ""
+            display_schema_keys = ",".join(
+                f"'{key}'" for key in saved_staged__schemas_m2m.keys()
             )
             logger.save(
-                f"saved {len(saved_feature_sets)} feature set{s} for slot{s}:"
-                f" {display_feature_set_keys}"
+                f"saved {len(saved_staged__schemas_m2m)} feature set{s} for slot{s}:"
+                f" {display_schema_keys}"
             )
 
 
-def save_feature_set_links(self: Artifact | Collection) -> None:
+def save_schema_links(self: Artifact | Collection) -> None:
     from lamindb._save import bulk_create
 
     Data = self.__class__
-    if hasattr(self, "_feature_sets"):
+    if hasattr(self, "_staged__schemas_m2m"):
         links = []
         host_id_field = get_host_id_field(self)
-        for slot, feature_set in self._feature_sets.items():
+        for slot, schema in self._staged__schemas_m2m.items():
             kwargs = {
                 host_id_field: self.id,
-                "featureset_id": feature_set.id,
+                "schema_id": schema.id,
                 "slot": slot,
             }
-            links.append(Data.feature_sets.through(**kwargs))
+            links.append(Data._schemas_m2m.through(**kwargs))
         bulk_create(links, ignore_conflicts=True)
 
 
@@ -140,7 +135,7 @@ def _describe_postgres(self: Artifact | Collection, print_types: bool = False):
             include_feature_link=True,
             include_fk=True,
             include_m2m=True,
-            include_featureset=True,
+            include_schema=True,
         )
     else:
         result = get_artifact_with_related(self, include_fk=True, include_m2m=True)
@@ -153,7 +148,7 @@ def _describe_postgres(self: Artifact | Collection, print_types: bool = False):
         tree=tree,
         related_data=related_data,
         with_labels=True,
-        print_params=hasattr(self, "type") and self.type == "model",
+        print_params=hasattr(self, "kind") and self.kind == "model",
     )
 
 
@@ -187,7 +182,7 @@ def _describe_sqlite(self: Artifact | Collection, print_types: bool = False):
         if isinstance(self, (Collection, Artifact)):
             many_to_many_fields.append("input_of_runs")
         if isinstance(self, Artifact):
-            many_to_many_fields.append("feature_sets")
+            many_to_many_fields.append("_schemas_m2m")
         self = (
             self.__class__.objects.using(self._state.db)
             .prefetch_related(*many_to_many_fields)
@@ -198,7 +193,7 @@ def _describe_sqlite(self: Artifact | Collection, print_types: bool = False):
         self,
         tree=tree,
         with_labels=True,
-        print_params=hasattr(self, "type") and self.type == "model",
+        print_params=hasattr(self, "kind") and self.kind == "kind",
     )
 
 
@@ -220,7 +215,7 @@ def validate_feature(feature: Feature, records: list[Record]) -> None:
     if not isinstance(feature, Feature):
         raise TypeError("feature has to be of type Feature")
     if feature._state.adding:
-        registries = {record.__class__.__get_name_with_schema__() for record in records}
+        registries = {record.__class__.__get_name_with_module__() for record in records}
         registries_str = "|".join(registries)
         msg = f"ln.Feature(name='{feature.name}', type='cat[{registries_str}]').save()"
         raise ValidationError(f"Feature not validated. If it looks correct: {msg}")
@@ -303,7 +298,7 @@ def add_labels(
                 " feature=ln.Feature(name='my_feature'))"
             )
         if feature.dtype.startswith("cat["):
-            orm_dict = dict_schema_name_to_model_name(Artifact)
+            orm_dict = dict_module_name_to_model_name(Artifact)
             for reg in feature.dtype.replace("cat[", "").rstrip("]").split("|"):
                 registry = orm_dict.get(reg)
                 records_validated += registry.from_values(records, field=field)
@@ -329,7 +324,7 @@ def add_labels(
         # strategy: group records by registry to reduce number of transactions
         records_by_related_name: dict = {}
         for record in records:
-            related_name = d.get(record.__class__.__get_name_with_schema__())
+            related_name = d.get(record.__class__.__get_name_with_module__())
             if related_name is None:
                 raise ValueError(f"Can't add labels to {record.__class__} record!")
             if related_name not in records_by_related_name:
@@ -340,15 +335,15 @@ def add_labels(
     else:
         validate_feature(feature, records)  # type:ignore
         records_by_registry = defaultdict(list)
-        feature_sets = self.feature_sets.filter(registry="Feature").all()
+        _schemas_m2m = self._schemas_m2m.filter(itype="Feature").all()
         internal_features = set()  # type: ignore
-        if len(feature_sets) > 0:
-            for feature_set in feature_sets:
+        if len(_schemas_m2m) > 0:
+            for schema in _schemas_m2m:
                 internal_features = internal_features.union(
-                    set(feature_set.members.values_list("name", flat=True))
+                    set(schema.members.values_list("name", flat=True))
                 )  # type: ignore
         for record in records:
-            records_by_registry[record.__class__.__get_name_with_schema__()].append(
+            records_by_registry[record.__class__.__get_name_with_module__()].append(
                 record
             )
         for registry_name, records in records_by_registry.items():
@@ -388,7 +383,7 @@ def add_labels(
 
 def _track_run_input(
     data: Artifact | Collection | Iterable[Artifact] | Iterable[Collection],
-    is_run_input: bool | None = None,
+    is_run_input: bool | Run | None = None,
     run: Run | None = None,
 ):
     # this is an internal hack right now for project-flow, but we can allow this

lamindb/core/_describe.py

@@ -7,18 +7,25 @@ from lamin_utils import logger
 from rich.text import Text
 from rich.tree import Tree
 
+from ._context import is_run_from_ipython
+
 if TYPE_CHECKING:
-    from lnschema_core.models import Artifact, Collection, Run
+    from lamindb.models import Artifact, Collection, Run
 
 
 def highlight_time(iso: str):
-    tz = datetime.datetime.now().astimezone().tzinfo
-    res = (
-        datetime.datetime.fromisoformat(iso)
-        .replace(tzinfo=datetime.timezone.utc)
-        .astimezone(tz)
-        .strftime("%Y-%m-%d %H:%M:%S")
-    )
+    try:
+        tz = datetime.datetime.now().astimezone().tzinfo
+        res = (
+            datetime.datetime.fromisoformat(iso)
+            .replace(tzinfo=datetime.timezone.utc)
+            .astimezone(tz)
+            .strftime("%Y-%m-%d %H:%M:%S")
+        )
+    except ValueError:
+        # raises ValueError: Invalid isoformat string: '<django.db.models.expressions.DatabaseDefault object at 0x1128ac440>'
+        # but can't be caught with `isinstance(iso, DatabaseDefault)` for unkown reasons
+        return Text("timestamp of unsaved record not available", style="dim")
     return Text(res, style="dim")
 
 
@@ -31,28 +38,33 @@ VALUES_WIDTH = 40
 def print_rich_tree(tree: Tree, fallback=str):
     from rich.console import Console
 
+    # If tree has no children, return fallback
+    if not tree.children:
+        return fallback
+
     console = Console(force_terminal=True)
+    printed = False
 
-    if tree.children:
-        try:
+    try:
+        if not is_run_from_ipython:
             from IPython import get_ipython
             from IPython.core.interactiveshell import InteractiveShell
             from IPython.display import display
 
             shell = get_ipython()
-            if isinstance(shell, InteractiveShell):  # Covers all interactive shells
+            if isinstance(shell, InteractiveShell):
                 display(tree)
+                printed = True
                 return ""
-            else:
-                with console.capture() as capture:
-                    console.print(tree)
-                return capture.get()
-        except (ImportError, NameError):
-            with console.capture() as capture:
-                console.print(tree)
-            return capture.get()
-    else:
-        return fallback
+    except (NameError, ImportError):
+        pass
+
+    # If not printed through IPython
+    if not printed:
+        # be careful to test this on a terminal
+        console = Console(force_terminal=True)
+        console.print(tree)
+        return ""
 
 
 def describe_header(self: Artifact | Collection | Run) -> Tree:
@@ -60,14 +72,14 @@ def describe_header(self: Artifact | Collection | Run) -> Tree:
         logger.warning(
             f"This is not the latest version of the {self.__class__.__name__}."
         )
-    if hasattr(self, "visibility"):
-        if self.visibility == 0:
+    if hasattr(self, "_branch_code"):
+        if self._branch_code == 0:
             logger.warning("This artifact is hidden.")
-        elif self.visibility == -1:
+        elif self._branch_code == -1:
             logger.warning("This artifact is the trash.")
     # initialize tree
     suffix = self.suffix if hasattr(self, "suffix") and self.suffix else ""
-    accessor = self._accessor if hasattr(self, "_accessor") and self._accessor else ""
+    accessor = self.otype if hasattr(self, "otype") and self.otype else ""
     suffix_accessor = (
         f"{suffix}/{accessor}" if suffix and accessor else suffix or accessor or ""
     )
@@ -89,15 +101,13 @@ def describe_general(self: Artifact | Collection, tree: Tree | None = None) -> T
     general = tree.add(Text("General", style="bold bright_cyan"))
     general.add(f".uid = '{self.uid}'")
     if hasattr(self, "key") and self.key:
-        general.add(
-            f".key = '{self.key}'" if self._key_is_virtual else f".key = {self.key}"
-        )
+        general.add(f".key = '{self.key}'")
     if hasattr(self, "size") and self.size:
         general.add(f".size = {self.size}")
     if hasattr(self, "hash") and self.hash:
         general.add(f".hash = '{self.hash}'")
-    if hasattr(self, "n_objects") and self.n_objects:
-        general.add(f".n_objects = {self.n_objects}")
+    if hasattr(self, "n_files") and self.n_files:
+        general.add(f".n_files = {self.n_files}")
     if hasattr(self, "n_observations") and self.n_observations:
         general.add(Text(f".n_observations = {self.n_observations}"))
     if hasattr(self, "version") and self.version:
@@ -131,7 +141,7 @@ def describe_general(self: Artifact | Collection, tree: Tree | None = None) -> T
     if hasattr(self, "transform") and self.transform:
         general.add(
             Text(
-                f".transform = '{self.transform.name}'",
+                f".transform = '{self.transform.description}'",
                 style="cyan3",
             )
         )

lamindb/core/_django.py

@@ -1,14 +1,17 @@
 from __future__ import annotations
 
+from functools import reduce
+
 from django.contrib.postgres.aggregates import ArrayAgg
 from django.db import connection
 from django.db.models import F, OuterRef, Q, Subquery
 from django.db.models.fields.related import ForeignKey, ManyToManyField
 from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel
 from django.db.models.functions import JSONObject
-from lnschema_core.models import Artifact, FeatureSet, Record
 
-from .schema import dict_related_model_to_related_name, get_schemas_modules
+from lamindb.models import Artifact, Record, Schema
+
+from .relations import dict_related_model_to_related_name, get_schema_modules
 
 
 def get_related_model(model, field_name):
@@ -32,7 +35,7 @@ def get_artifact_with_related(
     include_fk: bool = False,
     include_m2m: bool = False,
     include_feature_link: bool = False,
-    include_featureset: bool = False,
+    include_schema: bool = False,
 ) -> dict:
     """Fetch an artifact with its related data."""
     from lamindb._can_curate import get_name_field
@@ -40,12 +43,12 @@ def get_artifact_with_related(
     from ._label_manager import EXCLUDE_LABELS
 
     model = artifact.__class__
-    schema_modules = get_schemas_modules(artifact._state.db)
+    schema_modules = get_schema_modules(artifact._state.db)
 
     foreign_key_fields = [
         f.name
         for f in model._meta.fields
-        if f.is_relation and f.related_model.__get_schema_name__() in schema_modules
+        if f.is_relation and f.related_model.__get_module_name__() in schema_modules
     ]
 
     m2m_relations = (
@@ -81,15 +84,6 @@ def get_artifact_with_related(
                 id=F(f"{fk}__id"), name=F(f"{fk}__{name_field}")
             )
 
-    for name in m2m_relations:
-        related_model = get_related_model(model, name)
-        name_field = get_name_field(related_model)
-        annotations[f"m2mfield_{name}"] = ArrayAgg(
-            JSONObject(id=F(f"{name}__id"), name=F(f"{name}__{name_field}")),
-            filter=Q(**{f"{name}__isnull": False}),
-            distinct=True,
-        )
-
     for link in link_tables:
         link_model = getattr(model, link).rel.related_model
         if not hasattr(link_model, "feature"):
@@ -109,14 +103,14 @@ def get_artifact_with_related(
             .values("json_agg")
         )
 
-    if include_featureset:
-        annotations["featuresets"] = Subquery(
-            model.feature_sets.through.objects.filter(artifact=OuterRef("pk"))
+    if include_schema:
+        annotations["schemas"] = Subquery(
+            model._schemas_m2m.through.objects.filter(artifact=OuterRef("pk"))
             .annotate(
                 data=JSONObject(
                     id=F("id"),
                     slot=F("slot"),
-                    featureset=F("featureset"),
+                    schema=F("schema"),
                 )
             )
             .values("artifact")
@@ -135,25 +129,45 @@ def get_artifact_with_related(
     if not artifact_meta:
         return None
 
-    related_data: dict = {"m2m": {}, "fk": {}, "link": {}, "featuresets": {}}
+    related_data: dict = {"m2m": {}, "fk": {}, "link": {}, "schemas": {}}
     for k, v in artifact_meta.items():
-        if k.startswith("m2mfield_"):
-            related_data["m2m"][k[9:]] = v
-        elif k.startswith("fkfield_"):
+        if k.startswith("fkfield_"):
             related_data["fk"][k[8:]] = v
         elif k.startswith("linkfield_"):
             related_data["link"][k[10:]] = v
-        elif k == "featuresets":
+        elif k == "schemas":
             if v:
-                related_data["featuresets"] = get_featureset_m2m_relations(
-                    artifact, {i["featureset"]: i["slot"] for i in v}
+                related_data["schemas"] = get_schema_m2m_relations(
+                    artifact, {i["schema"]: i["slot"] for i in v}
                 )
 
-    related_data["m2m"] = {
-        k: {item["id"]: item["name"] for item in v}
-        for k, v in related_data["m2m"].items()
-        if v
-    }
+    if len(m2m_relations) == 0:
+        m2m_any = False
+    else:
+        m2m_any_expr = reduce(
+            lambda a, b: a | b,
+            (Q(**{f"{m2m_name}__isnull": False}) for m2m_name in m2m_relations),
+        )
+        # this is needed to avoid querying all m2m relations even if they are all empty
+        # this checks if non-empty m2m relations are present in the record
+        m2m_any = (
+            model.objects.using(artifact._state.db)
+            .filter(uid=artifact.uid)
+            .filter(m2m_any_expr)
+            .exists()
+        )
+    if m2m_any:
+        m2m_data = related_data["m2m"]
+        for m2m_name in m2m_relations:
+            related_model = get_related_model(model, m2m_name)
+            name_field = get_name_field(related_model)
+            m2m_records = (
+                getattr(artifact, m2m_name).values_list("id", name_field).distinct()
+            )
+            for rec_id, rec_name in m2m_records:
+                if m2m_name not in m2m_data:
+                    m2m_data[m2m_name] = {}
+                m2m_data[m2m_name][rec_id] = rec_name
 
     return {
         **{name: artifact_meta[name] for name in ["id", "uid"]},
@@ -161,30 +175,31 @@ def get_artifact_with_related(
     }
 
 
-def get_featureset_m2m_relations(
-    artifact: Artifact, slot_featureset: dict, limit: int = 20
-):
+def get_schema_m2m_relations(artifact: Artifact, slot_schema: dict, limit: int = 20):
     """Fetch all many-to-many relationships for given feature sets."""
     from lamindb._can_curate import get_name_field
 
     m2m_relations = [
         v
-        for v in dict_related_model_to_related_name(FeatureSet).values()
-        if not v.startswith("_") and v != "artifacts"
+        for v in dict_related_model_to_related_name(Schema).values()
+        if v is not None and not v.startswith("_") and v != "artifacts"
     ]
 
     annotations = {}
     related_names = {}
     for name in m2m_relations:
-        related_model = get_related_model(FeatureSet, name)
+        related_model = get_related_model(Schema, name)
+        if related_model is Schema:
+            # this is for the `type` field
+            continue
         name_field = get_name_field(related_model)
 
         # Get the correct field names for the through table
-        through_model = getattr(FeatureSet, name).through
+        through_model = getattr(Schema, name).through
 
         # Subquery to get limited related records
         limited_related = Subquery(
-            through_model.objects.filter(featureset=OuterRef("pk")).values(
+            through_model.objects.filter(schema=OuterRef("pk")).values(
                 related_model.__name__.lower()
             )[:limit]
         )
@@ -198,18 +213,18 @@ def get_featureset_m2m_relations(
             ),
             distinct=True,
         )
-        related_names[name] = related_model.__get_name_with_schema__()
+        related_names[name] = related_model.__get_name_with_module__()
 
-    featureset_m2m = (
-        FeatureSet.objects.using(artifact._state.db)
-        .filter(id__in=slot_featureset.keys())
+    schema_m2m = (
+        Schema.objects.using(artifact._state.db)
+        .filter(id__in=slot_schema.keys())
         .annotate(**annotations)
         .values("id", *annotations.keys())
     )
 
     result = {}
-    for fs in featureset_m2m:
-        slot = slot_featureset.get(fs["id"])
+    for fs in schema_m2m:
+        slot = slot_schema.get(fs["id"])
         result[fs["id"]] = (
             slot,
             {