text: from typing import Optional, Union
  
  import numpy as np
  from anndata import AnnData
  from pandas.api.types import CategoricalDtype
  
  from scvi.data._utils import _make_obs_column_categorical
  
  from ._obs_field import CategoricalObsField
  
  
  class LabelsWithUnlabeledObsField(CategoricalObsField):
      """
      An AnnDataField for labels which include explicitly unlabeled cells.
  
      Remaps the unlabeled category to the final index if present in labels.
      The unlabeled category is a specific category name specified by the user.
  
      Parameters
      ----------
      registry_key
          Key to register field under in data registry.
      obs_key
          Key to access the field in the AnnData obs mapping. If None, defaults to `registry_key`.
      unlabeled_category
          Value assigned to unlabeled cells.
      """
  
      UNLABELED_CATEGORY = "unlabeled_category"
      WAS_REMAPPED = "was_remapped"
  
      def __init__(
          self,
          registry_key: str,
          obs_key: Optional[str],
          unlabeled_category: Union[str, int, float],
      ) -> None:
          super().__init__(registry_key, obs_key)
          self._unlabeled_category = unlabeled_category
  
      def _remap_unlabeled_to_final_category(
          self, adata: AnnData, mapping: np.ndarray
      ) -> dict:
          labels = self._get_original_column(adata)
  
          if self._unlabeled_category in labels:
              unlabeled_idx = np.where(mapping == self._unlabeled_category)
              unlabeled_idx = unlabeled_idx[0][0]
              # move unlabeled category to be the last position
              mapping[unlabeled_idx], mapping[-1] = mapping[-1], mapping[unlabeled_idx]
          # could be in mapping in transfer case
          elif self._unlabeled_category not in mapping:
              # just put as last category
              mapping = np.asarray(list(mapping) + [self._unlabeled_category])
  
          cat_dtype = CategoricalDtype(categories=mapping, ordered=True)
          # rerun setup for the batch column
          mapping = _make_obs_column_categorical(
              adata,
              self._original_attr_key,
              self.attr_key,
              categorical_dtype=cat_dtype,
          )
  
          return {
              self.CATEGORICAL_MAPPING_KEY: mapping,
              self.ORIGINAL_ATTR_KEY: self._original_attr_key,
              self.UNLABELED_CATEGORY: self._unlabeled_category,
          }
  
      def register_field(self, adata: AnnData) -> dict:
          if self.is_default:
              self._setup_default_attr(adata)
  
          state_registry = super().register_field(adata)
          mapping = state_registry[self.CATEGORICAL_MAPPING_KEY]
          return self._remap_unlabeled_to_final_category(adata, mapping)
  
      def transfer_field(
          self,
          state_registry: dict,
          adata_target: AnnData,
          extend_categories: bool = False,
          **kwargs,
      ) -> dict:
          transfer_state_registry = super().transfer_field(
              state_registry, adata_target, extend_categories=extend_categories, **kwargs
          )
          mapping = transfer_state_registry[self.CATEGORICAL_MAPPING_KEY]
          return self._remap_unlabeled_to_final_category(adata_target, mapping)
  