Adding initial batch of samtools type hinted wrapper functions

fgpyo/sam/__init__.py

@@ -197,13 +197,14 @@ class SamFileType(enum.Enum):
         ext (str): The standard file extension for this file type.
     """
 
-    def __init__(self, mode: str, ext: str) -> None:
+    def __init__(self, mode: str, ext: str, index_ext: Optional[str]) -> None:
         self.mode = mode
         self.extension = ext
+        self.index_extension = index_ext
 
-    SAM = ("", ".sam")
-    BAM = ("b", ".bam")
-    CRAM = ("c", ".cram")
+    SAM = ("", ".sam", None)
+    BAM = ("b", ".bam", ".bai")
+    CRAM = ("c", ".cram", ".crai")
 
     @classmethod
     def from_path(cls, path: Union[Path, str]) -> "SamFileType":

fgpyo/sam/builder.py

@@ -27,6 +27,7 @@
 from pysam import AlignmentHeader
 
 from fgpyo import sam
+from fgpyo import samtools
 from fgpyo.sam import SamOrder
 
 
@@ -544,7 +545,7 @@ def to_path(
 
         with NamedTemporaryFile(suffix=".bam", delete=True) as fp:
             file_handle: IO
-            if self.sort_order is SamOrder.Unsorted:
+            if self.sort_order in [SamOrder.Unsorted, SamOrder.Unknown]:
                 file_handle = path.open("w")
             else:
                 file_handle = fp.file
@@ -555,20 +556,14 @@ def to_path(
                 for rec in self._records:
                     if pred(rec):
                         writer.write(rec)
-
-            default_samtools_opt_list = ["-o", str(path), fp.name]
-
             file_handle.close()
-            if self.sort_order == SamOrder.QueryName:
-                # Ignore type hints for now until we have wrappers to use here.
-                pysam.sort(*(["-n"] + default_samtools_opt_list))  # type: ignore
-            elif self.sort_order == SamOrder.Coordinate:
-                # Ignore type hints for now until we have wrappers to use here.
-                pysam.sort(*default_samtools_opt_list)  # type: ignore
-                if index:
-                    # Ignore type hints for now until we have wrappers to use here.
-                    pysam.index(str(path))  # type: ignore
-
+            if self.sort_order not in [SamOrder.Unsorted, SamOrder.Unknown]:
+                samtools.sort(
+                    input=Path(fp.name),
+                    output=path,
+                    index_output=index,
+                    sort_order=self.sort_order,
+                )
         return path
 
     def __len__(self) -> int: