fixup

structlib/atoms.py

@@ -96,6 +96,16 @@ def _selection_to_expr(selection: AtomSelection) -> polars.Expr:
     return _values_to_expr(selection, ty=polars.Boolean)
 
 
+def _selection_to_numpy(df: t.Union[polars.DataFrame, HasAtoms], selection: AtomSelection) -> NDArray[numpy.bool_]:
+    series = _selection_to_series(df, selection)
+    try:
+        import pyarrow
+    except ImportError:
+        # workaround without pyarrow
+        return numpy.array(series.to_list(), dtype=numpy.bool_)
+    return series.to_numpy(zero_copy_only=False)
+
+
 def _select_schema(df: t.Union[polars.DataFrame, HasAtoms], schema: SchemaDict) -> polars.DataFrame:
     """
     Select columns from ``self`` and cast to the given schema.
@@ -593,7 +603,7 @@ def with_symbol(self: HasAtomsT, symbols: ArrayLike, selection: t.Optional[AtomS
         Return `self` with the given atomic symbols.
         """
         if selection is not None:
-            selection = _selection_to_series(self, selection).to_numpy()
+            selection = _selection_to_numpy(self, selection)
             new_symbols = self.get_column('symbol')
             new_symbols[selection] = polars.Series(list(numpy.broadcast_to(symbols, len(selection))), dtype=polars.Utf8)
             symbols = new_symbols
@@ -607,7 +617,7 @@ def with_coords(self: HasAtomsT, pts: ArrayLike, selection: t.Optional[AtomSelec
         Return `self` replaced with the given atomic positions.
         """
         if selection is not None:
-            selection = _selection_to_series(self, selection).to_numpy()
+            selection = _selection_to_numpy(self, selection)
             new_pts = self.coords()
             pts = numpy.atleast_2d(pts)
             assert pts.shape[-1] == 3
@@ -778,4 +788,4 @@ def _repr_pretty_(self, p, cycle: bool) -> None:
 
 __all__ = [
     'Atoms', 'HasAtoms', 'IntoAtoms', 'AtomSelection', 'AtomValues',
-]
+]