ENH: Generalization of RangeIndex to support floats. (pandas-dev#46484)

Author: kapiliyer

pandas/core/dtypes/common.py

@@ -112,33 +112,34 @@ def ensure_str(value: bytes | Any) -> str:
     return value
 
 
-def ensure_python_int(value: int | np.integer) -> int:
+def ensure_python_intfloat(
+    value: int | float | np.integer | np.inexact, coerce_float: bool
+) -> int | float:
     """
-    Ensure that a value is a python int.
+    Ensure that a value is a python int or float.
 
     Parameters
     ----------
-    value: int or numpy.integer
+    value: int or float or numpy.integer or numpy.complex_
+    coerce_float: bool
 
     Returns
     -------
-    int
+    int or float
 
     Raises
     ------
-    TypeError: if the value isn't an int or can't be converted to one.
+    TypeError: if the value isn't an int/float and can't be converted to either.
     """
-    if not (is_integer(value) or is_float(value)):
+    if is_integer(value) or is_float(value):
+        new_value = float(value) if coerce_float else int(value)
+        assert np.isclose(value, new_value) if coerce_float else value == new_value
+    else:
         if not is_scalar(value):
             raise TypeError(
                 f"Value needs to be a scalar value, was type {type(value).__name__}"
             )
         raise TypeError(f"Wrong type {type(value)} for value {value}")
-    try:
-        new_value = int(value)
-        assert new_value == value
-    except (TypeError, ValueError, AssertionError) as err:
-        raise TypeError(f"Wrong type {type(value)} for value {value}") from err
     return new_value
 
 
@@ -1822,7 +1823,7 @@ def is_all_strings(value: ArrayLike) -> bool:
     "DT64NS_DTYPE",
     "ensure_float",
     "ensure_float64",
-    "ensure_python_int",
+    "ensure_python_intfloat",
     "ensure_str",
     "get_dtype",
     "infer_dtype_from_object",

pandas/core/indexes/float_range.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+class float_range:
+    """
+    Modified Python built-in function range that accepts floats as arguments.
+
+    Parameters
+    ----------
+    start : float (default: 0.0)
+        If "stop" is not given, interpreted as "stop" instead.
+    stop : float (default: 0.0)
+    step : float (default: 1.0)
+
+    Attributes
+    ----------
+    start
+    stop
+    step
+    current
+    """
+
+    def __init__(self, start=0.0, stop=0.0, step=1.0):
+        if stop is None:
+            stop, start = start, 0.0
+        self.start = start
+        self.current = start - step
+        self.stop = stop
+        self.step = step
+
+    def __contains__(self, key):
+        closeness_to_range = (key - self.start) % self.step
+        return (
+            (self.start < key or np.isclose(self.start, key))
+            and (key < self.stop and not np.isclose(self.stop, key))
+            and (
+                np.isclose(closeness_to_range, 0)
+                or np.isclose(closeness_to_range, self.step)
+            )
+        )
+
+    def __eq__(self, other):
+        return (
+            np.isclose(self.start, other.start)
+            and np.isclose(self.stop, other.stop)
+            and np.isclose(self.step, other.step)
+        )
+
+    def __getitem__(self, key):
+        # Allows for reversing
+        if isinstance(key, slice):
+            if key.start is None and key.stop is None and key.step == -1:
+                return float_range(
+                    start=self.stop - self.step,
+                    stop=self.start - self.step,
+                    step=-self.step,
+                )
+            raise IndexError
+        raise TypeError
+
+    def __iter__(self):
+        return self
+
+    def __len__(self):
+        return self.length
+
+    def __next__(self):
+        self.current += self.step
+        if self.current < self.stop:
+            return self.current
+        raise StopIteration
+
+    def index(self, key):
+        if key not in self:
+            raise ValueError
+        return round((key - self.start) / self.step)
+
+    @property
+    def length(self):
+        if self.step > 0:
+            lo, hi = self.start, self.stop
+            step = self.step
+        else:
+            hi, lo = self.start, self.stop
+            step = -self.step
+        if lo > hi or np.isclose(lo, hi):
+            return 0
+        else:
+            return round((hi - lo - step) / step + 1)

pandas/core/indexes/range.py

@@ -35,11 +35,11 @@
 
 from pandas.core.dtypes.common import (
     ensure_platform_int,
-    ensure_python_int,
+    ensure_python_intfloat,
     is_float,
     is_integer,
+    is_numeric_dtype,
     is_scalar,
-    is_signed_integer_dtype,
     is_timedelta64_dtype,
 )
 from pandas.core.dtypes.generic import ABCTimedeltaIndex
@@ -50,6 +50,7 @@
 from pandas.core.construction import extract_array
 import pandas.core.indexes.base as ibase
 from pandas.core.indexes.base import maybe_extract_name
+from pandas.core.indexes.float_range import float_range
 from pandas.core.indexes.numeric import (
     Float64Index,
     Int64Index,
@@ -67,20 +68,20 @@ class RangeIndex(NumericIndex):
     """
     Immutable Index implementing a monotonic integer range.
 
-    RangeIndex is a memory-saving special case of Int64Index limited to
-    representing monotonic ranges. Using RangeIndex may in some instances
-    improve computing speed.
+    RangeIndex is a memory-saving special case of NumericIndex
+    limited to representing monotonic ranges. Using RangeIndex may in some
+    instances improve computing speed.
 
-    This is the default index type used
-    by DataFrame and Series when no explicit index is provided by the user.
+    This is the default index type used by DataFrame and Series when no
+    explicit index is provided by the user.
 
     Parameters
     ----------
-    start : int (default: 0), range, or other RangeIndex instance
-        If int and "stop" is not given, interpreted as "stop" instead.
-    stop : int (default: 0)
-    step : int (default: 1)
-    dtype : np.int64
+    start : int|float (default: 0), range, or other RangeIndex instance
+        If int|float and "stop" is not given, interpreted as "stop" instead.
+    stop : int|float (default: 0)
+    step : int|float (default: 1)
+    dtype : np.int64|np.float64
         Unused, accepted for homogeneity with other index types.
     copy : bool, default False
         Unused, accepted for homogeneity with other index types.
@@ -101,16 +102,17 @@ class RangeIndex(NumericIndex):
     --------
     Index : The base pandas Index type.
     Int64Index : Index of int64 data.
+    Float64Index : Index of float64 data
     """
 
     _typ = "rangeindex"
-    _dtype_validation_metadata = (is_signed_integer_dtype, "signed integer")
-    _range: range
+    _dtype_validation_metadata = (is_numeric_dtype, "numeric type")
+    _range: range | float_range
     _is_backward_compat_public_numeric_index: bool = False
 
     @property
-    def _engine_type(self) -> type[libindex.Int64Engine]:
-        return libindex.Int64Engine
+    def _engine_type(self) -> type[libindex.IndexEngine]:
+        return libindex.Float64Engine if is_float(self.start) else libindex.Int64Engine
 
     # --------------------------------------------------------------------
     # Constructors
@@ -135,25 +137,36 @@ def __new__(
 
         # validate the arguments
         if com.all_none(start, stop, step):
-            raise TypeError("RangeIndex(...) must be called with integers")
+            raise TypeError("RangeIndex(...) must be called with integers/floats")
 
-        start = ensure_python_int(start) if start is not None else 0
+        # check if any of the arguments is a float
+        coerce_float = any(is_float(attr) for attr in [start, stop, step])
+
+        start = ensure_python_intfloat(start, coerce_float) if start is not None else 0
 
         if stop is None:
-            start, stop = 0, start
+            start, stop = 0.0 if coerce_float else 0, start
         else:
-            stop = ensure_python_int(stop)
-
-        step = ensure_python_int(step) if step is not None else 1
+            stop = ensure_python_intfloat(stop, coerce_float)
+
+        step = (
+            ensure_python_intfloat(step, coerce_float)
+            if step is not None
+            else 1.0
+            if coerce_float
+            else 1
+        )
         if step == 0:
             raise ValueError("Step must not be zero")
 
-        rng = range(start, stop, step)
+        rng = (
+            float_range(start, stop, step) if coerce_float else range(start, stop, step)
+        )
         return cls._simple_new(rng, name=name)
 
     @classmethod
     def from_range(
-        cls, data: range, name=None, dtype: Dtype | None = None
+        cls, data: range | float_range, name=None, dtype: Dtype | None = None
     ) -> RangeIndex:
         """
         Create RangeIndex from a range object.
@@ -162,7 +175,7 @@ def from_range(
         -------
         RangeIndex
         """
-        if not isinstance(data, range):
+        if not isinstance(data, range) and not isinstance(data, float_range):
             raise TypeError(
                 f"{cls.__name__}(...) must be called with object coercible to a "
                 f"range, {repr(data)} was passed"
@@ -171,10 +184,12 @@ def from_range(
         return cls._simple_new(data, name=name)
 
     @classmethod
-    def _simple_new(cls, values: range, name: Hashable = None) -> RangeIndex:
+    def _simple_new(
+        cls, values: range | float_range, name: Hashable = None
+    ) -> RangeIndex:
         result = object.__new__(cls)
 
-        assert isinstance(values, range)
+        assert isinstance(values, range) or isinstance(values, float_range)
 
         result._range = values
         result._name = name
@@ -184,12 +199,12 @@ def _simple_new(cls, values: range, name: Hashable = None) -> RangeIndex:
 
     # --------------------------------------------------------------------
 
-    # error: Return type "Type[Int64Index]" of "_constructor" incompatible with return
+    # error: Return type "Type[NumericIndex]" of "_constructor" incompatible with return
     # type "Type[RangeIndex]" in supertype "Index"
     @cache_readonly
-    def _constructor(self) -> type[Int64Index]:  # type: ignore[override]
+    def _constructor(self) -> type[NumericIndex]:  # type: ignore[override]
         """return the class to use for construction"""
-        return Int64Index
+        return Float64Index if is_float(self.start) else Int64Index
 
     # error: Signature of "_data" incompatible with supertype "Index"
     @cache_readonly
@@ -199,7 +214,12 @@ def _data(self) -> np.ndarray:  # type: ignore[override]
 
         The constructed array is saved in ``_cache``.
         """
-        return np.arange(self.start, self.stop, self.step, dtype=np.int64)
+        return np.arange(
+            self.start,
+            self.stop,
+            self.step,
+            dtype=np.float64 if is_float(self.start) else np.int64,
+        )
 
     def _get_data_as_items(self):
         """return a list of tuples of start, stop, step"""
@@ -352,7 +372,7 @@ def memory_usage(self, deep: bool = False) -> int:
 
     @property
     def dtype(self) -> np.dtype:
-        return np.dtype(np.int64)
+        return np.dtype(np.float64 if is_float(self.start) else np.int64)
 
     @property
     def is_unique(self) -> bool:
@@ -370,14 +390,14 @@ def is_monotonic_decreasing(self) -> bool:
     def __contains__(self, key: Any) -> bool:
         hash(key)
         try:
-            key = ensure_python_int(key)
+            key = ensure_python_intfloat(key, True)
         except TypeError:
             return False
         return key in self._range
 
     @property
     def inferred_type(self) -> str:
-        return "integer"
+        return "float" if is_float(self.start) else "integer"
 
     # --------------------------------------------------------------------
     # Indexing Methods

pandas/tests/indexes/ranges/test_constructors.py

@@ -153,15 +153,6 @@ def test_constructor_corner(self):
 
         tm.assert_index_equal(index, expected, exact="equiv")
 
-        # non-int raise Exception
+        # non-int/float raise Exception
         with pytest.raises(TypeError, match=r"Wrong type \<class 'str'\>"):
             RangeIndex("1", "10", "1")
-        with pytest.raises(TypeError, match=r"Wrong type \<class 'float'\>"):
-            RangeIndex(1.1, 10.2, 1.3)
-
-        # invalid passed type
-        with pytest.raises(
-            ValueError,
-            match="Incorrect `dtype` passed: expected signed integer, received float64",
-        ):
-            RangeIndex(1, 5, dtype="float64")