actually fix indexing

This commit is contained in:
sneakers-the-rat 2024-08-07 19:22:29 -07:00
parent 362965daf5
commit cebb21993d
Signed by untrusted user who does not match committer: jonny
GPG key ID: 6DCB96EF1E4D232D
14 changed files with 134 additions and 148 deletions

View file

@ -109,7 +109,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -127,10 +127,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -145,14 +147,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -258,7 +258,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -266,6 +266,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -286,7 +287,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -304,10 +305,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -322,14 +325,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -258,7 +258,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -266,6 +266,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -286,7 +287,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -304,10 +305,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -322,14 +325,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -258,7 +258,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -266,6 +266,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -286,7 +287,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -304,10 +305,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -322,14 +325,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -259,7 +259,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -267,6 +267,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -287,7 +288,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -305,10 +306,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -323,14 +326,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -259,7 +259,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -267,6 +267,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -287,7 +288,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -305,10 +306,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -323,14 +326,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -259,7 +259,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -267,6 +267,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -287,7 +288,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -305,10 +306,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -323,14 +326,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -259,7 +259,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -267,6 +267,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -287,7 +288,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -305,10 +306,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -323,14 +326,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -259,7 +259,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -267,6 +267,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -287,7 +288,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -305,10 +306,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -323,14 +326,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -259,7 +259,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -267,6 +267,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -287,7 +288,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -305,10 +306,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -323,14 +326,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -259,7 +259,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -267,6 +267,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -287,7 +288,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -305,10 +306,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -323,14 +326,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -259,7 +259,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -267,6 +267,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -287,7 +288,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -305,10 +306,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -323,14 +326,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -259,7 +259,7 @@ class DynamicTableMixin(BaseModel):
]: ... ]: ...
@overload @overload
def __getitem__(self, item: slice) -> DataFrame: ... def __getitem__(self, item: Union[slice, "NDArray"]) -> DataFrame: ...
def __getitem__( def __getitem__(
self, self,
@ -267,6 +267,7 @@ class DynamicTableMixin(BaseModel):
str, str,
int, int,
slice, slice,
"NDArray",
Tuple[int, Union[int, str]], Tuple[int, Union[int, str]],
Tuple[Union[int, slice], ...], Tuple[Union[int, slice], ...],
], ],
@ -287,7 +288,7 @@ class DynamicTableMixin(BaseModel):
if isinstance(item, str): if isinstance(item, str):
return self._columns[item] return self._columns[item]
if isinstance(item, (int, slice, np.integer, np.ndarray)): if isinstance(item, (int, slice, np.integer, np.ndarray)):
return DataFrame.from_dict(self._slice_range(item)) data = self._slice_range(item)
elif isinstance(item, tuple): elif isinstance(item, tuple):
if len(item) != 2: if len(item) != 2:
raise ValueError( raise ValueError(
@ -305,10 +306,12 @@ class DynamicTableMixin(BaseModel):
return self._columns[cols][rows] return self._columns[cols][rows]
data = self._slice_range(rows, cols) data = self._slice_range(rows, cols)
return DataFrame.from_dict(data)
else: else:
raise ValueError(f"Unsure how to get item with key {item}") raise ValueError(f"Unsure how to get item with key {item}")
# cast to DF
return DataFrame(data)
def _slice_range( def _slice_range(
self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None self, rows: Union[int, slice, np.ndarray], cols: Optional[Union[str, List[str]]] = None
) -> Dict[str, Union[list, "NDArray", "VectorData"]]: ) -> Dict[str, Union[list, "NDArray", "VectorData"]]:
@ -323,14 +326,10 @@ class DynamicTableMixin(BaseModel):
else: else:
val = self._columns[k][rows] val = self._columns[k][rows]
if isinstance(val, BaseModel): # scalars need to be wrapped in series for pandas
# special case where pandas will unpack a pydantic model if not isinstance(rows, (Iterable, slice)):
# into {n_fields} rows, rather than keeping it in a dict
val = Series([val])
elif isinstance(rows, int) and hasattr(val, "shape") and val.shape and len(val) > 1:
# special case where we are returning a row in a ragged array,
# same as above - prevent pandas pivoting to long
val = Series([val]) val = Series([val])
data[k] = val data[k] = val
return data return data

View file

@ -150,6 +150,24 @@ def test_dynamictable_indexing(electrical_series):
assert subsection.dtypes.values.tolist() == dtypes[0:3] assert subsection.dtypes.values.tolist() == dtypes[0:3]
def test_dynamictable_ragged(units):
"""
Should be able to index ragged arrays using an implicit _index column
Also tests:
- passing arrays directly instead of wrapping in vectordata/index specifically,
if the models in the fixture instantiate then this works
"""
units, spike_times, spike_idx = units
# ensure we don't pivot to long when indexing
assert units[0].shape[0] == 1
# check that we got the indexing boundaries corrunect
# (and that we are forwarding attr calls to the dataframe by accessing shape
for i in range(units.shape[0]):
assert np.all(units.iloc[i, 0] == spike_times[i])
def test_dynamictable_region_basic(electrical_series): def test_dynamictable_region_basic(electrical_series):
""" """
DynamicTableRegion should be able to refer to a row or rows of another table DynamicTableRegion should be able to refer to a row or rows of another table
@ -175,7 +193,7 @@ def test_dynamictable_region_basic(electrical_series):
# getting a list of table rows is actually correct behavior here because # getting a list of table rows is actually correct behavior here because
# this list of table rows is actually the cell of another table # this list of table rows is actually the cell of another table
rows = series.electrodes[0:3] rows = series.electrodes[0:3]
assert all([row.id == idx for row, idx in zip(rows, [4, 3, 2])]) assert all([all(row.id == idx) for row, idx in zip(rows, [4, 3, 2])])
def test_dynamictable_region_ragged(): def test_dynamictable_region_ragged():
@ -218,24 +236,6 @@ def test_dynamictable_region_ragged():
assert all([all(row[1].timeseries == i) for i, row in zip([1, 2, 3], rows.iterrows())]) assert all([all(row[1].timeseries == i) for i, row in zip([1, 2, 3], rows.iterrows())])
def test_dynamictable_ragged(units):
"""
Should be able to index ragged arrays using an implicit _index column
Also tests:
- passing arrays directly instead of wrapping in vectordata/index specifically,
if the models in the fixture instantiate then this works
"""
units, spike_times, spike_idx = units
# ensure we don't pivot to long when indexing
assert units[0].shape[0] == 1
# check that we got the indexing boundaries corrunect
# (and that we are forwarding attr calls to the dataframe by accessing shape
for i in range(units.shape[0]):
assert np.all(units.iloc[i, 0] == spike_times[i])
def test_dynamictable_append_column(): def test_dynamictable_append_column():
pass pass