@@ -729,3 +729,54 @@ def test_restore_dtype_on_multiindexes(dtype: str) -> None:
729
729
foo = xr .Dataset (coords = {"bar" : ("bar" , np .array ([0 , 1 ], dtype = dtype ))})
730
730
foo = foo .stack (baz = ("bar" ,))
731
731
assert str (foo ["bar" ].values .dtype ) == dtype
732
+
733
+
734
+ class IndexWithExtraVariables (Index ):
735
+ @classmethod
736
+ def from_variables (cls , variables , * , options = None ):
737
+ return cls ()
738
+
739
+ def create_variables (self , variables = None ):
740
+ if variables is None :
741
+ # For Coordinates.from_xindex(), return all variables the index can create
742
+ return {
743
+ "time" : Variable (dims = ("time" ,), data = [1 , 2 , 3 ]),
744
+ "valid_time" : Variable (
745
+ dims = ("time" ,),
746
+ data = [2 , 3 , 4 ], # time + 1
747
+ attrs = {"description" : "time + 1" },
748
+ ),
749
+ }
750
+
751
+ result = dict (variables )
752
+ if "time" in variables :
753
+ result ["valid_time" ] = Variable (
754
+ dims = ("time" ,),
755
+ data = variables ["time" ].data + 1 ,
756
+ attrs = {"description" : "time + 1" },
757
+ )
758
+ return result
759
+
760
+
761
+ def test_set_xindex_with_extra_variables () -> None :
762
+ """Test that set_xindex raises an error when custom index creates extra variables."""
763
+
764
+ ds = xr .Dataset (coords = {"time" : [1 , 2 , 3 ]}).reset_index ("time" )
765
+
766
+ # Test that set_xindex raises error for extra variables
767
+ with pytest .raises (ValueError , match = "extra variables 'valid_time'" ):
768
+ ds .set_xindex ("time" , IndexWithExtraVariables )
769
+
770
+
771
+ def test_set_xindex_factory_method_pattern () -> None :
772
+ ds = xr .Dataset (coords = {"time" : [1 , 2 , 3 ]}).reset_index ("time" )
773
+
774
+ # Test the recommended factory method pattern
775
+ coord_vars = {"time" : ds ._variables ["time" ]}
776
+ index = IndexWithExtraVariables .from_variables (coord_vars )
777
+ coords = xr .Coordinates .from_xindex (index )
778
+ result = ds .assign_coords (coords )
779
+
780
+ assert "time" in result .variables
781
+ assert "valid_time" in result .variables
782
+ assert_array_equal (result .valid_time .data , result .time .data + 1 )
0 commit comments