From 77b42636336b45085f6d93c4f781949f0965c747 Mon Sep 17 00:00:00 2001
From: Charles Tapley Hoyt <cthoyt@gmail.com>
Date: Mon, 6 Nov 2023 20:56:24 +0100
Subject: [PATCH] Update type hints for prefix map merging (#454)

This PR does the following:

1. Creates a literal type for the 3 different prefix map merge modes
2. Updates type hints based on that
---
 src/sssom/cli.py       | 10 +++++-----
 src/sssom/constants.py | 14 +++++---------
 src/sssom/io.py        | 12 ++++++++----
 3 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/src/sssom/cli.py b/src/sssom/cli.py
index ae4f23b9..61266f25 100644
--- a/src/sssom/cli.py
+++ b/src/sssom/cli.py
@@ -16,7 +16,7 @@
 import sys
 from operator import itemgetter
 from pathlib import Path
-from typing import Any, Callable, List, Optional, TextIO, Tuple
+from typing import Any, Callable, List, Optional, TextIO, Tuple, get_args
 
 import click
 import curies
@@ -28,7 +28,7 @@
 
 from sssom.constants import (
     DEFAULT_VALIDATION_TYPES,
-    PREFIX_MAP_MODES,
+    MergeMode,
     SchemaValidationType,
     _get_sssom_schema_object,
 )
@@ -172,9 +172,9 @@ def convert(input: str, output: TextIO, output_format: str):
     default="metadata_only",
     show_default=True,
     required=True,
-    type=click.Choice(PREFIX_MAP_MODES, case_sensitive=False),
+    type=click.Choice(get_args(MergeMode), case_sensitive=False),
     help="Defines whether the prefix map in the metadata should be extended or replaced with "
-    "the SSSOM default prefix map. Must be one of metadata_only, sssom_default_only, merged",
+    "the SSSOM default prefix map.",
 )
 @click.option(
     "-p",
@@ -205,7 +205,7 @@ def parse(
     input: str,
     input_format: str,
     metadata: str,
-    prefix_map_mode: str,
+    prefix_map_mode: MergeMode,
     clean_prefixes: bool,
     strict_clean_prefixes: bool,
     output: TextIO,
diff --git a/src/sssom/constants.py b/src/sssom/constants.py
index 7b4f8e42..84ecd18b 100644
--- a/src/sssom/constants.py
+++ b/src/sssom/constants.py
@@ -4,7 +4,7 @@
 import uuid
 from enum import Enum
 from functools import lru_cache
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Literal
 
 import pkg_resources
 import yaml
@@ -38,14 +38,10 @@
 ]
 
 UNKNOWN_IRI = "http://w3id.org/sssom/unknown_prefix/"
-PREFIX_MAP_MODE_METADATA_ONLY = "metadata_only"
-PREFIX_MAP_MODE_SSSOM_DEFAULT_ONLY = "sssom_default_only"
-PREFIX_MAP_MODE_MERGED = "merged"
-PREFIX_MAP_MODES = [
-    PREFIX_MAP_MODE_METADATA_ONLY,
-    PREFIX_MAP_MODE_SSSOM_DEFAULT_ONLY,
-    PREFIX_MAP_MODE_MERGED,
-]
+MergeMode = Literal["metadata_only", "sssom_default_only", "merged"]
+PREFIX_MAP_MODE_METADATA_ONLY: MergeMode = "metadata_only"
+PREFIX_MAP_MODE_SSSOM_DEFAULT_ONLY: MergeMode = "sssom_default_only"
+PREFIX_MAP_MODE_MERGED: MergeMode = "merged"
 ENTITY_REFERENCE = "EntityReference"
 
 # Slot Constants
diff --git a/src/sssom/io.py b/src/sssom/io.py
index 28f9eaf8..7f16e39c 100644
--- a/src/sssom/io.py
+++ b/src/sssom/io.py
@@ -21,6 +21,7 @@
     PREFIX_MAP_MODE_MERGED,
     PREFIX_MAP_MODE_METADATA_ONLY,
     PREFIX_MAP_MODE_SSSOM_DEFAULT_ONLY,
+    MergeMode,
     MetadataType,
     SchemaValidationType,
     get_default_metadata,
@@ -52,9 +53,10 @@ def convert_file(
 def parse_file(
     input_path: str,
     output: TextIO,
+    *,
     input_format: Optional[str] = None,
     metadata_path: Optional[str] = None,
-    prefix_map_mode: Optional[str] = None,
+    prefix_map_mode: Optional[MergeMode] = None,
     clean_prefixes: bool = True,
     strict_clean_prefixes: bool = True,
     embedded_mode: bool = True,
@@ -68,7 +70,7 @@ def parse_file(
     :param metadata_path: The path to a file containing the sssom metadata (including prefix_map)
         to be used during parse.
     :param prefix_map_mode: Defines whether the prefix map in the metadata should be extended or replaced with
-        the SSSOM default prefix map. Must be one of metadata_only, sssom_default_only, merged
+        the SSSOM default prefix map derived from the :mod:`bioregistry`.
     :param clean_prefixes: If True (default), records with unknown prefixes are removed from the SSSOM file.
     :param strict_clean_prefixes: If True (default), clean_prefixes() will be in strict mode.
     :param embedded_mode:If True (default), the dataframe and metadata are exported in one file (tsv), else two separate files (tsv and yaml).
@@ -129,7 +131,7 @@ def split_file(input_path: str, output_directory: Union[str, Path]) -> None:
 
 
 def get_metadata_and_prefix_map(
-    metadata_path: Union[None, str, Path] = None, *, prefix_map_mode: Optional[str] = None
+    metadata_path: Union[None, str, Path] = None, *, prefix_map_mode: Optional[MergeMode] = None
 ) -> Tuple[Converter, MetadataType]:
     """
     Load SSSOM metadata from a YAML file, and then augment it with default prefixes.
@@ -150,7 +152,9 @@ def get_metadata_and_prefix_map(
     return converter, metadata
 
 
-def _merge_converter(converter: Converter, prefix_map_mode: str = None) -> Converter:
+def _merge_converter(
+    converter: Converter, prefix_map_mode: Optional[MergeMode] = None
+) -> Converter:
     """Merge the metadata's converter with the default converter."""
     if prefix_map_mode is None or prefix_map_mode == PREFIX_MAP_MODE_METADATA_ONLY:
         return converter