Split PyRanges construction from pandas' internal frame fallback

.gitignore

@@ -2,6 +2,16 @@ docs/chipseq.csv
 docs/chipseq.gtf
 docs/test.bed
 
+# files the doctests write into the repo root when run from there
+/chipseq.csv
+/chipseq.gtf
+/test.bed
+/outfile.fasta
+/temp*.fasta
+/temp*.fasta.fai
+/Dgyro.taa_CDS.gtf
+/Dgyro_taa_CDS_seqs.tsv
+
 *.swp
 .snakemake/
 wheelhouse/

pyranges1/core/pyranges_main.py

@@ -2,7 +2,7 @@
 
 import logging
 import sys
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional, cast
 
@@ -168,23 +168,28 @@ class PyRanges(RangeFrame):
 
     """
 
-    def __new__(cls, *args, **kwargs) -> "pr.PyRanges | pd.DataFrame":  # type: ignore[misc]
+    def __new__(cls, *args, **kwargs) -> "pr.PyRanges":  # type: ignore[misc]
         """Create a new instance of a PyRanges object."""
         # __new__ is a special static method used for creating and
         # returning a new instance of a class. It is called before
         # __init__ and is typically used in scenarios requiring
         # control over the creation of new instances
 
-        # Logic to decide whether to return an instance of PyRanges or a DataFrame
+        # Direct construction always yields a real PyRanges, or raises. The
+        # graceful fallback to a plain DataFrame (e.g. when a pandas operation
+        # like .drop() removes a required column) lives in _constructor instead,
+        # so it only kicks in for pandas' internal frame reconstruction and never
+        # for a user calling PyRanges(...) directly. See geopandas.GeoDataFrame
+        # for the same split between __init__ and _constructor.
         if not args and "data" not in kwargs:
-            df = pd.DataFrame({k: [] for k in GENOME_LOC_COLS})
-            df.__class__ = pr.PyRanges
-            return df
+            return super().__new__(cls)
 
-        df = pd.DataFrame(kwargs.get("../data") or args[0])
+        df = pd.DataFrame(kwargs.get("data") if "data" in kwargs else args[0])
         missing_any_required_columns = not set(GENOME_LOC_COLS).issubset({*df.columns})
         if missing_any_required_columns:
-            return df
+            missing = sorted(set(GENOME_LOC_COLS) - set(df.columns))
+            msg = f"Cannot construct PyRanges: missing required column(s) {missing}."
+            raise ValueError(msg)
 
         return super().__new__(cls)
 
@@ -202,8 +207,21 @@ def __init__(self, *args, **kwargs) -> None:
         self._loci = LociGetter(self)
 
     @property
-    def _constructor(self) -> type:
-        return pr.PyRanges
+    def _constructor(self) -> Callable[..., "pr.PyRanges | pd.DataFrame"]:
+        return self._constructor_with_fallback
+
+    @classmethod
+    def _constructor_with_fallback(cls, *args, **kwargs) -> "pr.PyRanges | pd.DataFrame":
+        """Build a PyRanges, falling back to a DataFrame if a required column is missing.
+
+        Used by pandas internally to reconstruct frames from operations (e.g. .drop(),
+        groupby aggregations). See the "Operations that remove a column required for a
+        PyRanges return a DataFrame instead" example above.
+        """
+        df = pd.DataFrame(*args, **kwargs)
+        if not set(GENOME_LOC_COLS).issubset({*df.columns}):
+            return df
+        return cls(df)
 
     def groupby(self, *args, **kwargs) -> "PyRangesDataFrameGroupBy":
         """Groupby PyRanges."""
@@ -4852,7 +4870,8 @@ def remove_strand(self) -> "PyRanges":
         """
         if not self.has_strand:
             return self
-        return self.drop_and_return(STRAND_COL, axis=1)
+        # Strand isn't a required column, so this drop can never fall back to a DataFrame.
+        return cast("PyRanges", self.drop_and_return(STRAND_COL, axis=1))
 
     def flip_strand(self: "PyRanges") -> "PyRanges":
         """Flip the strand of every interval (+ → - and - → +).

pyranges1/ext/orfs.py

@@ -2,7 +2,7 @@
 import sys
 import warnings
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 import pandas as pd
@@ -126,7 +126,8 @@ def calculate_frame(p: "pr.PyRanges", group_by: str | list[str], frame_col: str
 
     gr[frame_col] = sorted_p[FRAME_COL]
 
-    return gr.drop_and_return(TEMP_INDEX_COL, axis=1)
+    # TEMP_INDEX_COL isn't a required column, so this drop can never fall back to a DataFrame.
+    return cast("pr.PyRanges", gr.drop_and_return(TEMP_INDEX_COL, axis=1))
 
 
 def extend_orfs(  # noqa: C901,PLR0912,PLR0915
@@ -410,7 +411,8 @@ def pverbose(msg: Any) -> None:
             stacklevel=2,
         )
 
-    p = p.drop_and_return(["__length"], axis=1)
+    # "__length" isn't a required column, so this drop can never fall back to a DataFrame.
+    p = cast("pr.PyRanges", p.drop_and_return(["__length"], axis=1))
     ##################
 
     # Load Sequence Data from a Fasta file

pyranges1/range_frame/range_frame.py

@@ -1,7 +1,7 @@
 import inspect
 import warnings
 from collections.abc import Callable, Iterable
-from typing import Any, TypeVar
+from typing import Any, TypeVar, cast
 
 import numpy as np
 import pandas as pd
@@ -220,7 +220,9 @@ def combine_interval_columns(
 
         cols_to_drop = list({start, end, start2, end2}.difference(RANGE_COLS) if drop_old_columns else {})
 
-        return z.drop_and_return(labels=cols_to_drop, axis="columns")
+        # cols_to_drop excludes RANGE_COLS (Start/End), so this drop can never remove a
+        # required column and fall back to a plain DataFrame.
+        return cast("RangeFrame", z.drop_and_return(labels=cols_to_drop, axis="columns"))
 
     def cluster_overlaps(
         self,
@@ -856,15 +858,27 @@ def ensure_valid_ranges(self) -> "RangeFrame":
     def copy(self, *args, **kwargs) -> "RangeFrame":  # pyright: ignore[reportIncompatibleMethodOverride]  # noqa: D102
         return _mypy_ensure_rangeframe(super().copy(*args, **kwargs))
 
-    def drop(self, *args, **kwargs) -> "RangeFrame | None":  # type: ignore[override]  # noqa: D102
-        return self.__class__(super().drop(*args, **kwargs))
+    @classmethod
+    def _constructor_with_fallback(cls, *args, **kwargs) -> "RangeFrame | pd.DataFrame":
+        """Build a cls, falling back to a plain DataFrame if cls cannot represent the data.
 
-    def drop_and_return[T: "RangeFrame"](self: T, *args: Any, **kwargs: Any) -> T:  # noqa: PYI019, D102
+        A RangeFrame has no required columns, so it never falls back; PyRanges overrides
+        this to degrade to a DataFrame when a required column is missing.
+
+        This is not pandas' _constructor property: that one is called by pandas whenever it
+        rebuilds a frame internally, while this one is only called by the methods below.
+        """
+        return cls(*args, **kwargs)
+
+    def drop(self, *args, **kwargs) -> "RangeFrame | pd.DataFrame | None":  # type: ignore[override]  # noqa: D102
+        return self._constructor_with_fallback(super().drop(*args, **kwargs))
+
+    def drop_and_return(self, *args: Any, **kwargs: Any) -> "RangeFrame | pd.DataFrame":  # noqa: D102
         kwargs["inplace"] = False
-        return self.__class__(super().drop(*args, **kwargs))
+        return self._constructor_with_fallback(super().drop(*args, **kwargs))
 
-    def reindex(self, *args, **kwargs) -> "RangeFrame":  # noqa: D102
-        return self.__class__(super().reindex(*args, **kwargs))
+    def reindex(self, *args, **kwargs) -> "RangeFrame | pd.DataFrame":  # noqa: D102
+        return self._constructor_with_fallback(super().reindex(*args, **kwargs))
 
 
 def _mypy_ensure_rangeframe(r: pd.DataFrame) -> "RangeFrame":

tests/unit/test_pandas_overrides.py

@@ -236,3 +236,28 @@ def test_groupby_getattr_series_as_index_false(gr) -> None:
     res = result.agg("first")
     # DataFrame because as_index=False
     assert isinstance(res, pd.DataFrame)
+
+
+def test_drop(gr) -> None:
+    assert type(gr.drop("Val", axis=1)) is pr.PyRanges
+
+    # dropping a column a PyRanges requires degrades to a DataFrame
+    assert type(gr.drop("Chromosome", axis=1)) is pd.DataFrame
+
+
+def test_drop_and_return(gr) -> None:
+    assert type(gr.drop_and_return("Val", axis=1)) is pr.PyRanges
+    assert type(gr.drop_and_return("Chromosome", axis=1)) is pd.DataFrame
+
+
+def test_reindex(gr) -> None:
+    assert type(gr.reindex(columns=["Chromosome", "Start", "End"])) is pr.PyRanges
+    assert type(gr.reindex(columns=["Start", "End"])) is pd.DataFrame
+
+
+def test_range_frame_never_degrades() -> None:
+    # a RangeFrame requires no columns, so these always stay a RangeFrame
+    rf = pr.RangeFrame({"Start": [0, 10], "End": [40, 20], "Val": [50, 30]})
+    assert type(rf.drop("Val", axis=1)) is pr.RangeFrame
+    assert type(rf.drop_and_return("Val", axis=1)) is pr.RangeFrame
+    assert type(rf.reindex(columns=["Start"])) is pr.RangeFrame