lamindb 1.2a2__py3-none-any.whl → 1.3.1__py3-none-any.whl

lamindb/models/feature.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import importlib
 from typing import TYPE_CHECKING, Any, get_args, overload
 
+import numpy as np
 import pandas as pd
 from django.db import models
 from django.db.models import CASCADE, PROTECT, Q
@@ -12,6 +13,7 @@ from lamin_utils import logger
 from lamindb_setup._init_instance import get_schema_module_name
 from lamindb_setup.core.hashing import HASH_LENGTH, hash_dict
 from pandas.api.types import CategoricalDtype, is_string_dtype
+from pandas.core.dtypes.base import ExtensionDtype
 
 from lamindb.base.fields import (
     BooleanField,
@@ -20,7 +22,7 @@ from lamindb.base.fields import (
     JSONField,
     TextField,
 )
-from lamindb.base.types import FeatureDtype, FieldAttr
+from lamindb.base.types import Dtype, FieldAttr
 from lamindb.errors import FieldValidationError, ValidationError
 
 from ..base.ids import base62_12
@@ -36,18 +38,43 @@ from .run import (
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
-    from pandas.core.dtypes.base import ExtensionDtype
-
     from .schema import Schema
 
-FEATURE_DTYPES = set(get_args(FeatureDtype))
+FEATURE_DTYPES = set(get_args(Dtype))
+
 
+def parse_dtype(dtype_str: str, is_param: bool = False) -> list[dict[str, str]]:
+    """Parses feature data type string into a structured list of components."""
+    from .artifact import Artifact
 
-def parse_dtype_single_cat(
+    allowed_dtypes = FEATURE_DTYPES
+    if is_param:
+        allowed_dtypes.add("dict")
+    is_composed_cat = dtype_str.startswith("cat[") and dtype_str.endswith("]")
+    result = []
+    if is_composed_cat:
+        related_registries = dict_module_name_to_model_name(Artifact)
+        registries_str = dtype_str.replace("cat[", "")[:-1]  # strip last ]
+        if registries_str != "":
+            registry_str_list = registries_str.split("|")
+            for cat_single_dtype_str in registry_str_list:
+                single_result = parse_cat_dtype(
+                    cat_single_dtype_str, related_registries
+                )
+                result.append(single_result)
+    elif dtype_str not in allowed_dtypes:
+        raise ValueError(
+            f"dtype is '{dtype_str}' but has to be one of {FEATURE_DTYPES}!"
+        )
+    return result
+
+
+def parse_cat_dtype(
     dtype_str: str,
     related_registries: dict[str, Record] | None = None,
     is_itype: bool = False,
-) -> dict:
+) -> dict[str, Any]:
+    """Parses a categorical dtype string into its components (registry, field, subtypes)."""
     from .artifact import Artifact
 
     assert isinstance(dtype_str, str)  # noqa: S101
@@ -115,38 +142,18 @@ def parse_dtype_single_cat(
     }
 
 
-def parse_dtype(dtype_str: str, is_param: bool = False) -> list[dict[str, str]]:
-    from .artifact import Artifact
-
-    allowed_dtypes = FEATURE_DTYPES
-    if is_param:
-        allowed_dtypes.add("dict")
-    is_composed_cat = dtype_str.startswith("cat[") and dtype_str.endswith("]")
-    result = []
-    if is_composed_cat:
-        related_registries = dict_module_name_to_model_name(Artifact)
-        registries_str = dtype_str.replace("cat[", "")[:-1]  # strip last ]
-        if registries_str != "":
-            registry_str_list = registries_str.split("|")
-            for cat_single_dtype_str in registry_str_list:
-                single_result = parse_dtype_single_cat(
-                    cat_single_dtype_str, related_registries
-                )
-                result.append(single_result)
-    elif dtype_str not in allowed_dtypes:
-        raise ValueError(
-            f"dtype is '{dtype_str}' but has to be one of {FEATURE_DTYPES}!"
-        )
-    return result
-
-
-def get_dtype_str_from_dtype(dtype: Any, is_itype: bool = False) -> str:
+def serialize_dtype(
+    dtype: Record | FieldAttr | list[Record], is_itype: bool = False
+) -> str:
+    """Converts a data type object into its string representation."""
     if (
         not isinstance(dtype, list)
         and hasattr(dtype, "__name__")
         and dtype.__name__ in FEATURE_DTYPES
     ):
         dtype_str = dtype.__name__
+    elif isinstance(dtype, (ExtensionDtype, np.dtype)):
+        dtype_str = serialize_pandas_dtype(dtype)
     else:
         error_message = (
             "dtype has to be a record, a record field, or a list of records, not {}"
@@ -177,7 +184,7 @@ def get_dtype_str_from_dtype(dtype: Any, is_itype: bool = False) -> str:
     return dtype_str
 
 
-def convert_pandas_dtype_to_lamin_dtype(pandas_dtype: ExtensionDtype) -> str:
+def serialize_pandas_dtype(pandas_dtype: ExtensionDtype) -> str:
     if is_string_dtype(pandas_dtype):
         if not isinstance(pandas_dtype, CategoricalDtype):
             dtype = "str"
@@ -189,6 +196,8 @@ def convert_pandas_dtype_to_lamin_dtype(pandas_dtype: ExtensionDtype) -> str:
     else:
         # strip precision qualifiers
         dtype = "".join(dt for dt in pandas_dtype.name if not dt.isdigit())
+        if dtype == "uint":
+            dtype = "int"
     if dtype.startswith("datetime"):
         dtype = dtype.split("[")[0]
     assert dtype in FEATURE_DTYPES  # noqa: S101
@@ -220,7 +229,7 @@ def process_init_feature_param(args, kwargs, is_param: bool = False):
     dtype_str = None
     if dtype is not None:
         if not isinstance(dtype, str):
-            dtype_str = get_dtype_str_from_dtype(dtype)
+            dtype_str = serialize_dtype(dtype)
         else:
             dtype_str = dtype
             parse_dtype(dtype_str, is_param=is_param)
@@ -247,14 +256,16 @@ class Feature(Record, CanCurate, TracksRun, TracksUpdates):
 
     Args:
         name: `str` Name of the feature, typically.  column name.
-        dtype: `FeatureDtype | Registry | list[Registry] | FieldAttr` See :class:`~lamindb.base.types.FeatureDtype`.
-            For categorical types, can define from which registry values are
-            sampled, e.g., `ULabel` or `[ULabel, bionty.CellType]`.
+        dtype: `Dtype | Registry | list[Registry] | FieldAttr` See :class:`~lamindb.base.types.Dtype`.
+            For categorical types, you can define to which registry values are
+            restricted, e.g., `ULabel` or `[ULabel, bionty.CellType]`.
         unit: `str | None = None` Unit of measure, ideally SI (`"m"`, `"s"`, `"kg"`, etc.) or `"normalized"` etc.
         description: `str | None = None` A description.
         synonyms: `str | None = None` Bar-separated synonyms.
         nullable: `bool = True` Whether the feature can have null-like values (`None`, `pd.NA`, `NaN`, etc.), see :attr:`~lamindb.Feature.nullable`.
         default_value: `Any | None = None` Default value for the feature.
+        coerce_dtype: `bool = False` When True, attempts to coerce values to the specified dtype
+            during validation, see :attr:`~lamindb.Feature.coerce_dtype`.
         cat_filters: `dict[str, str] | None = None` Subset a registry by additional filters to define valid categories.
 
     Note:
@@ -323,6 +334,7 @@ class Feature(Record, CanCurate, TracksRun, TracksUpdates):
     _aux_fields: dict[str, tuple[str, type]] = {
         "0": ("default_value", bool),
         "1": ("nullable", bool),
+        "2": ("coerce_dtype", bool),
     }
 
     id: int = models.AutoField(primary_key=True)
@@ -333,13 +345,8 @@ class Feature(Record, CanCurate, TracksRun, TracksUpdates):
     """Universal id, valid across DB instances."""
     name: str = CharField(max_length=150, db_index=True, unique=True)
     """Name of feature (hard unique constraint `unique=True`)."""
-    dtype: FeatureDtype | None = CharField(db_index=True, null=True)
-    """Data type (:class:`~lamindb.base.types.FeatureDtype`).
-
-    For categorical types, can define from which registry values are
-    sampled, e.g., `'cat[ULabel]'` or `'cat[bionty.CellType]'`. Unions are also
-    allowed if the feature samples from two registries, e.g., `'cat[ULabel|bionty.CellType]'`
-    """
+    dtype: Dtype | None = CharField(db_index=True, null=True)
+    """Data type (:class:`~lamindb.base.types.Dtype`)."""
     type: Feature | None = ForeignKey(
         "self", PROTECT, null=True, related_name="records"
     )
@@ -381,7 +388,7 @@ class Feature(Record, CanCurate, TracksRun, TracksUpdates):
 
     Is stored as a list rather than a tuple because it's serialized as JSON.
     """
-    proxy_dtype: FeatureDtype | None = CharField(default=None, null=True)
+    proxy_dtype: Dtype | None = CharField(default=None, null=True)
     """Proxy data type.
 
     If the feature is an image it's often stored via a path to the image file. Hence, while the dtype might be
@@ -411,7 +418,7 @@ class Feature(Record, CanCurate, TracksRun, TracksUpdates):
     def __init__(
         self,
         name: str,
-        dtype: FeatureDtype | Registry | list[Registry] | FieldAttr,
+        dtype: Dtype | Registry | list[Registry] | FieldAttr,
         type: Feature | None = None,
         is_type: bool = False,
         unit: str | None = None,
@@ -419,6 +426,7 @@ class Feature(Record, CanCurate, TracksRun, TracksUpdates):
         synonyms: str | None = None,
         nullable: bool = True,
         default_value: str | None = None,
+        coerce_dtype: bool = False,
         cat_filters: dict[str, str] | None = None,
     ): ...
 
@@ -440,10 +448,12 @@ class Feature(Record, CanCurate, TracksRun, TracksUpdates):
         default_value = kwargs.pop("default_value", None)
         nullable = kwargs.pop("nullable", True)  # default value of nullable
         cat_filters = kwargs.pop("cat_filters", None)
+        coerce_dtype = kwargs.pop("coerce_dtype", False)
         kwargs = process_init_feature_param(args, kwargs)
         super().__init__(*args, **kwargs)
         self.default_value = default_value
         self.nullable = nullable
+        self.coerce_dtype = coerce_dtype
         dtype_str = kwargs.pop("dtype", None)
         if cat_filters:
             assert "|" not in dtype_str  # noqa: S101
@@ -476,7 +486,7 @@ class Feature(Record, CanCurate, TracksRun, TracksUpdates):
             if name in categoricals:
                 dtypes[name] = "cat"
             else:
-                dtypes[name] = convert_pandas_dtype_to_lamin_dtype(col.dtype)
+                dtypes[name] = serialize_pandas_dtype(col.dtype)
         with logger.mute():  # silence the warning "loaded record with exact same name "
             features = [
                 Feature(name=name, dtype=dtype) for name, dtype in dtypes.items()
@@ -489,6 +499,25 @@ class Feature(Record, CanCurate, TracksRun, TracksUpdates):
         super().save(*args, **kwargs)
         return self
 
+    @property
+    def coerce_dtype(self) -> bool:
+        """Whether dtypes should be coerced during validation.
+
+        For example, a `objects`-dtyped pandas column can be coerced to `categorical` and would pass validation if this is true.
+        """
+        if self._aux is not None and "af" in self._aux and "2" in self._aux["af"]:  # type: ignore
+            return self._aux["af"]["2"]  # type: ignore
+        else:
+            return False
+
+    @coerce_dtype.setter
+    def coerce_dtype(self, value: bool) -> None:
+        if self._aux is None:  # type: ignore
+            self._aux = {}  # type: ignore
+        if "af" not in self._aux:
+            self._aux["af"] = {}
+        self._aux["af"]["2"] = value
+
     @property
     def default_value(self) -> Any:
         """A default value that overwrites missing values (default `None`).

lamindb/models/project.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
 from django.core.validators import RegexValidator
 from django.db import models
@@ -66,6 +66,23 @@ class Person(Record, CanCurate, TracksRun, TracksUpdates, ValidateFields):
     external: bool = BooleanField(default=True, db_index=True)
     """Whether the person is external to the organization."""
 
+    @overload
+    def __init__(
+        self,
+        name: str,
+        email: str | None = None,
+        external: bool = True,
+    ): ...
+
+    @overload
+    def __init__(
+        self,
+        *db_args,
+    ): ...
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
 
 class Reference(Record, CanCurate, TracksRun, TracksUpdates, ValidateFields):
     """References such as internal studies, papers, documents, or URLs.
@@ -94,12 +111,6 @@ class Reference(Record, CanCurate, TracksRun, TracksUpdates, ValidateFields):
     """Universal id, valid across DB instances."""
     name: str = CharField(db_index=True)
     """Title or name of the reference document."""
-    abbr: str | None = CharField(
-        max_length=32,
-        db_index=True,
-        null=True,
-    )
-    """An abbreviation for the reference."""
     type: Reference | None = ForeignKey(
         "self", PROTECT, null=True, related_name="records"
     )
@@ -111,6 +122,12 @@ class Reference(Record, CanCurate, TracksRun, TracksUpdates, ValidateFields):
     """Records of this type."""
     is_type: bool = BooleanField(default=False, db_index=True, null=True)
     """Distinguish types from instances of the type."""
+    abbr: str | None = CharField(
+        max_length=32,
+        db_index=True,
+        null=True,
+    )
+    """An abbreviation for the reference."""
     url: str | None = URLField(null=True)
     """URL linking to the reference."""
     pubmed_id: int | None = BigIntegerField(null=True, db_index=True)
@@ -147,6 +164,30 @@ class Reference(Record, CanCurate, TracksRun, TracksUpdates, ValidateFields):
     )
     """Collections associated with this reference."""
 
+    @overload
+    def __init__(
+        self,
+        name: str,
+        type: Reference | None = None,
+        is_type: bool = False,
+        abbr: str | None = None,
+        url: str | None = None,
+        pubmed_id: int | None = None,
+        doi: str | None = None,
+        description: str | None = None,
+        text: str | None = None,
+        date: DateType | None = None,
+    ): ...
+
+    @overload
+    def __init__(
+        self,
+        *db_args,
+    ): ...
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
 
 class Project(Record, CanCurate, TracksRun, TracksUpdates, ValidateFields):
     """Projects.
@@ -241,6 +282,27 @@ class Project(Record, CanCurate, TracksRun, TracksUpdates, ValidateFields):
     _status_code: int = models.SmallIntegerField(default=0, db_index=True)
     """Status code."""
 
+    @overload
+    def __init__(
+        self,
+        name: str,
+        type: Project | None = None,
+        is_type: bool = False,
+        abbr: str | None = None,
+        url: str | None = None,
+        start_date: DateType | None = None,
+        end_date: DateType | None = None,
+    ): ...
+
+    @overload
+    def __init__(
+        self,
+        *db_args,
+    ): ...
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
 
 class ArtifactProject(BasicRecord, LinkORM, TracksRun):
     id: int = models.BigAutoField(primary_key=True)

lamindb/models/query_set.py

@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Generic, NamedTuple, TypeVar, Union
 import pandas as pd
 from django.core.exceptions import FieldError
 from django.db import models
-from django.db.models import F, ForeignKey, ManyToManyField
+from django.db.models import F, ForeignKey, ManyToManyField, Subquery
 from django.db.models.fields.related import ForeignObjectRel
 from lamin_utils import logger
 from lamindb_setup.core._docs import doc_args
@@ -567,7 +567,17 @@ class QuerySet(models.QuerySet):
             include_kwargs = {s: F(s) for s in include if s not in field_names}
             annotate_kwargs.update(include_kwargs)
         if annotate_kwargs:
-            queryset = self.annotate(**annotate_kwargs)
+            id_subquery = self.values("id")
+            # for annotate, we want the queryset without filters so that joins don't affect the annotations
+            query_set_without_filters = self.model.objects.filter(
+                id__in=Subquery(id_subquery)
+            )
+            if self.query.order_by:
+                # Apply the same ordering to the new queryset
+                query_set_without_filters = query_set_without_filters.order_by(
+                    *self.query.order_by
+                )
+            queryset = query_set_without_filters.annotate(**annotate_kwargs)
         else:
             queryset = self
 

lamindb/models/record.py

@@ -13,7 +13,9 @@ from typing import (
     Any,
     Literal,
     NamedTuple,
+    TypeVar,
     Union,
+    overload,
 )
 
 import dj_database_url
@@ -49,6 +51,7 @@ from django.db.models.lookups import (
 )
 from lamin_utils import colors, logger
 from lamin_utils._lookup import Lookup
+from lamindb_setup import settings as setup_settings
 from lamindb_setup._connect_instance import (
     get_owner_name_from_identifier,
     load_instance_settings,
@@ -87,6 +90,7 @@ if TYPE_CHECKING:
     from .transform import Transform
 
 
+T = TypeVar("T", bound="Record")
 IPYTHON = getattr(builtins, "__IPYTHON__", False)
 
 
@@ -202,12 +206,12 @@ def validate_literal_fields(record: Record, kwargs) -> None:
         return None
     if record.__class__.__name__ in "Feature":
         return None
-    from lamindb.base.types import FeatureDtype, TransformType
+    from lamindb.base.types import Dtype, TransformType
 
     types = {
         "TransformType": TransformType,
-        "ArtifactKind": FeatureDtype,
-        "FeatureDtype": FeatureDtype,
+        "ArtifactKind": Dtype,
+        "Dtype": Dtype,
     }
     errors = {}
     annotations = getattr(record.__class__, "__annotations__", {})
@@ -466,19 +470,16 @@ class Registry(ModelBase):
         return QuerySet(model=cls, using=_using_key).filter(*queries, **expressions)
 
     def get(
-        cls,
+        cls: type[T],
         idlike: int | str | None = None,
         **expressions,
-    ) -> Record:
+    ) -> T:
         """Get a single record.
 
         Args:
             idlike: Either a uid stub, uid or an integer id.
             expressions: Fields and values passed as Django query expressions.
 
-        Returns:
-            A record.
-
         Raises:
             :exc:`docs:lamindb.errors.DoesNotExist`: In case no matching record is found.
 
@@ -486,9 +487,10 @@ class Registry(ModelBase):
             - Guide: :doc:`docs:registries`
             - Django documentation: `Queries <https://docs.djangoproject.com/en/stable/topics/db/queries/>`__
 
-        Examples:
-            >>> ulabel = ln.ULabel.get("FvtpPJLJ")
-            >>> ulabel = ln.ULabel.get(name="my-label")
+        Examples::
+
+            ulabel = ln.ULabel.get("FvtpPJLJ")
+            ulabel = ln.ULabel.get(name="my-label")
         """
         from .query_set import QuerySet
 
@@ -594,7 +596,11 @@ class Registry(ModelBase):
 
         if instance is None:
             return QuerySet(model=cls, using=None)
+
         owner, name = get_owner_name_from_identifier(instance)
+        if f"{owner}/{name}" == setup_settings.instance.slug:
+            return QuerySet(model=cls, using=None)
+
         settings_file = instance_settings_file(name, owner)
         cache_filepath = (
             ln_setup.settings.cache_dir / f"instance--{owner}--{name}--uid.txt"
@@ -737,8 +743,7 @@ class BasicRecord(models.Model, metaclass=Registry):
             )
         else:
             super().__init__(*args)
-            _store_record_old_name(self)
-            _store_record_old_key(self)
+        track_current_key_and_name_values(self)
 
     def save(self, *args, **kwargs) -> Record:
         """Save.
@@ -806,8 +811,8 @@ class BasicRecord(models.Model, metaclass=Registry):
                     init_self_from_db(self, pre_existing_record)
                 else:
                     raise
-            _store_record_old_name(self)
-            _store_record_old_key(self)
+            # call the below in case a user makes more updates to the record
+            track_current_key_and_name_values(self)
         # perform transfer of many-to-many fields
         # only supported for Artifact and Collection records
         if db is not None and db != "default" and using_key is None:
@@ -875,7 +880,13 @@ class BasicRecord(models.Model, metaclass=Registry):
 
 
 class Space(BasicRecord):
-    """Spaces."""
+    """Spaces.
+
+    You can use spaces to restrict access to records within an instance.
+
+    All data in this registry is synced from `lamin.ai` to enable re-using spaces across instances.
+    There is no need to manually create records.
+    """
 
     id: int = models.SmallAutoField(primary_key=True)
     """Internal id, valid only in one DB instance."""
@@ -901,6 +912,26 @@ class Space(BasicRecord):
     )
     """Creator of run."""
 
+    @overload
+    def __init__(
+        self,
+        name: str,
+        description: str | None = None,
+    ): ...
+
+    @overload
+    def __init__(
+        self,
+        *db_args,
+    ): ...
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
 
 @doc_args(RECORD_REGISTRY_EXAMPLE)
 class Record(BasicRecord, metaclass=Registry):
@@ -989,8 +1020,8 @@ def _get_record_kwargs(record_class) -> list[tuple[str, str]]:
     pattern = r"@overload\s+def __init__\s*\(([\s\S]*?)\):\s*\.{3}"
     overloads = re.finditer(pattern, source)
 
-    for overload in overloads:
-        params_block = overload.group(1)
+    for single_overload in overloads:
+        params_block = single_overload.group(1)
         # This is an additional safety measure if the overloaded signature that we're
         # looking for is not at the top but a "db_args" constructor
         if "*db_args" in params_block:
@@ -1037,13 +1068,14 @@ def _search(
     field: StrField | list[StrField] | None = None,
     limit: int | None = 20,
     case_sensitive: bool = False,
-    using_key: str | None = None,
     truncate_string: bool = False,
 ) -> QuerySet:
     if string is None:
         raise ValueError("Cannot search for None value! Please pass a valid string.")
 
-    input_queryset = _queryset(cls, using_key=using_key)
+    input_queryset = (
+        cls.all() if isinstance(cls, (QuerySet, Manager)) else cls.objects.all()
+    )
     registry = input_queryset.model
     name_field = getattr(registry, "_name_field", "name")
     if field is None:
@@ -1152,7 +1184,7 @@ def _lookup(
     using_key: str | None = None,
 ) -> NamedTuple:
     """{}"""  # noqa: D415
-    queryset = _queryset(cls, using_key=using_key)
+    queryset = cls.all() if isinstance(cls, (QuerySet, Manager)) else cls.objects.all()
     field = get_name_field(registry=queryset.model, field=field)
 
     return Lookup(
@@ -1172,7 +1204,7 @@ def _lookup(
 def get_name_field(
     registry: type[Record] | QuerySet | Manager,
     *,
-    field: str | StrField | None = None,
+    field: StrField | None = None,
 ) -> str:
     """Get the 1st char or text field from the registry."""
     if isinstance(registry, (QuerySet, Manager)):
@@ -1212,16 +1244,6 @@ def get_name_field(
     return field
 
 
-def _queryset(cls: Record | QuerySet | Manager, using_key: str) -> QuerySet:
-    if isinstance(cls, (QuerySet, Manager)):
-        return cls.all()
-    elif using_key is None or using_key == "default":
-        return cls.objects.all()
-    else:
-        # using must be called on cls, otherwise the connection isn't found
-        return cls.using(using_key).all()
-
-
 def add_db_connection(db: str, using: str):
     db_config = dj_database_url.config(
         default=db, conn_max_age=600, conn_health_checks=True
@@ -1374,18 +1396,14 @@ def transfer_to_default_db(
     return None
 
 
-def _store_record_old_name(record: Record):
-    # writes the name to the _name attribute, so we can detect renaming upon save
-    if hasattr(record, "_name_field"):
-        record._old_name = getattr(record, record._name_field)
-
-
-def _store_record_old_key(record: Record):
-    from lamindb.models import Artifact, Transform
+def track_current_key_and_name_values(record: Record):
+    from lamindb.models import Artifact
 
-    # writes the key to the _old_key attribute, so we can detect key changes upon save
-    if isinstance(record, (Artifact, Transform)):
+    if isinstance(record, Artifact):
         record._old_key = record.key
+        record._old_suffix = record.suffix
+    elif hasattr(record, "_name_field"):
+        record._old_name = getattr(record, record._name_field)
 
 
 def check_name_change(record: Record):
@@ -1466,20 +1484,29 @@ def check_key_change(record: Union[Artifact, Transform]):
 
     if not isinstance(record, Artifact) or not hasattr(record, "_old_key"):
         return
+    if record._old_suffix != record.suffix:
+        raise InvalidArgument(
+            f"Changing the `.suffix` of an artifact is not allowed! You tried to change it from '{record._old_suffix}' to '{record.suffix}'."
+        )
 
-    old_key = record._old_key or ""
-    new_key = record.key or ""
+    old_key = record._old_key
+    new_key = record.key
 
     if old_key != new_key:
         if not record._key_is_virtual:
             raise InvalidArgument(
-                f"Changing a non-virtual key of an artifact is not allowed! Tried to change key from '{old_key}' to '{new_key}'."
+                f"Changing a non-virtual key of an artifact is not allowed! You tried to change it from '{old_key}' to '{new_key}'."
             )
-        old_key_suffix = (
-            record.suffix
-            if record.suffix
-            else extract_suffix_from_path(PurePosixPath(old_key), arg_name="key")
-        )
+        if old_key is not None:
+            old_key_suffix = extract_suffix_from_path(
+                PurePosixPath(old_key), arg_name="key"
+            )
+            assert old_key_suffix == record.suffix, (  # noqa: S101
+                old_key_suffix,
+                record.suffix,
+            )
+        else:
+            old_key_suffix = record.suffix
         new_key_suffix = extract_suffix_from_path(
             PurePosixPath(new_key), arg_name="key"
         )