diff --git a/tiled/_tests/adapters/test_arrow.py b/tiled/_tests/adapters/test_arrow.py index 7515cb85b..24cf4b34a 100644 --- a/tiled/_tests/adapters/test_arrow.py +++ b/tiled/_tests/adapters/test_arrow.py @@ -4,7 +4,9 @@ import pytest from tiled.adapters.arrow import ArrowAdapter +from tiled.structures.core import StructureFamily from tiled.structures.table import TableStructure +from tiled.structures.data_source import DataSource, Management, Storage names = ["f0", "f1", "f2"] data0 = [ @@ -26,11 +28,30 @@ @pytest.fixture -def adapter() -> ArrowAdapter: +def data_source_from_init_storage() -> DataSource: table = pa.Table.from_arrays(data0, names) structure = TableStructure.from_arrow_table(table, npartitions=3) - assets = ArrowAdapter.init_storage(data_uri, structure=structure, path_parts=[]) - return ArrowAdapter([asset.data_uri for asset in assets], structure=structure) + data_source = DataSource( + management=Management.writable, + mimetype="application/vnd.apache.arrow.file", + structure_family=StructureFamily.table, + structure=structure, + assets=[], + ) + storage = Storage(filesystem=data_uri, sql=None) + return ArrowAdapter.init_storage( + data_source=data_source, storage=storage, path_parts=[] + ) + + +@pytest.fixture +def adapter(data_source_from_init_storage: DataSource) -> ArrowAdapter: + data_source = data_source_from_init_storage + return ArrowAdapter( + [asset.data_uri for asset in data_source.assets], + data_source.structure, + ) + def test_attributes(adapter: ArrowAdapter) -> None: diff --git a/tiled/_tests/test_catalog.py b/tiled/_tests/test_catalog.py index a4c915ec0..1af24c1f2 100644 --- a/tiled/_tests/test_catalog.py +++ b/tiled/_tests/test_catalog.py @@ -268,7 +268,7 @@ async def test_write_dataframe_external_direct(a, tmpdir): async def test_write_array_internal_direct(a, tmpdir): arr = numpy.ones((5, 3)) ad = ArrayAdapter.from_array(arr) - structure = asdict(ad.structure()) + structure = ad.structure() await a.create_node( key="x", structure_family="array", @@ -521,7 +521,7 @@ async def test_constraints_on_parameter_and_num(a, assets): DataSource( structure_family=arr_adapter.structure_family, mimetype="text/csv", - structure=asdict(arr_adapter.structure()), + structure=arr_adapter.structure(), parameters={}, management=Management.external, assets=assets, diff --git a/tiled/adapters/arrow.py b/tiled/adapters/arrow.py index 08ab7aa4c..eed744f79 100644 --- a/tiled/adapters/arrow.py +++ b/tiled/adapters/arrow.py @@ -70,9 +70,9 @@ def init_storage( Class to initialize the list of assets for given uri. Parameters ---------- - data_uri : - structure : - + storage : the storage option for .arrow files + data_source : data source representing the adapter + path_parts: the list of partitions Returns ------- The list of assets. @@ -81,6 +81,7 @@ def init_storage( data_uri = str(storage.get("filesystem")) + "".join( f"/{quote_plus(segment)}" for segment in path_parts ) + print('uri', data_uri) directory = path_from_uri(data_uri) directory.mkdir(parents=True, exist_ok=True) assets = [ diff --git a/tiled/adapters/sql.py b/tiled/adapters/sql.py index b57b5659b..b1cf5a93a 100644 --- a/tiled/adapters/sql.py +++ b/tiled/adapters/sql.py @@ -88,7 +88,7 @@ def init_storage( ---------- storage: the storage kind data_source : data source describing the data - path_parts: ? + path_parts: the list of partitions. Returns ------- A modified copy of the data source diff --git a/tiled/adapters/zarr.py b/tiled/adapters/zarr.py index b92e1d50e..3459606aa 100644 --- a/tiled/adapters/zarr.py +++ b/tiled/adapters/zarr.py @@ -13,9 +13,10 @@ from ..adapters.utils import IndexersMixin from ..iterviews import ItemsView, KeysView, ValuesView +from ..server.schemas import DataSource from ..structures.array import ArrayStructure from ..structures.core import Spec, StructureFamily -from ..structures.data_source import Asset, DataSource, Storage +from ..structures.data_source import Asset, Storage from ..type_aliases import JSON, NDSlice from ..utils import Conflicts, node_repr, path_from_uri from .array import ArrayAdapter, slice_and_shape_from_block_and_chunks diff --git a/tiled/catalog/adapter.py b/tiled/catalog/adapter.py index 93c7fd836..aebad928e 100644 --- a/tiled/catalog/adapter.py +++ b/tiled/catalog/adapter.py @@ -1192,6 +1192,10 @@ def _prepare_structure(structure_family, structure): "Convert from pydantic model to dict." if structure is None: return None + if isinstance(structure, dict): + return structure + if dataclasses.is_dataclass(structure): + return dataclasses.asdict(structure) return structure.model_dump() diff --git a/tiled/server/schemas.py b/tiled/server/schemas.py index 7c9dc3d66..67f449d02 100644 --- a/tiled/server/schemas.py +++ b/tiled/server/schemas.py @@ -149,10 +149,10 @@ def from_orm(cls, orm: tiled.catalog.orm.Revision) -> Revision: class DataSource(pydantic.BaseModel, Generic[StructureT]): - id: Optional[int] + id: Optional[int] = None structure_family: StructureFamily structure: Optional[StructureT] - mimetype: Optional[str] + mimetype: Optional[str] = None parameters: dict = {} assets: List[Asset] = [] management: Management = Management.writable