|
3 | 3 | import functools
|
4 | 4 | import logging
|
5 | 5 | import time
|
6 |
| -from dataclasses import dataclass |
7 | 6 | from datetime import timedelta
|
8 | 7 | from pathlib import Path
|
9 | 8 | from typing import (Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, Dict, List,
|
|
31 | 30 | stream_status_updates, stream_unique_updates_only)
|
32 | 31 |
|
33 | 32 |
|
34 |
| -@dataclass |
35 |
| -class StreamInfoUrl: |
36 |
| - '''Contains information on accessing ServiceX data via a url |
37 |
| - ''' |
38 |
| - url: str |
39 |
| - file: str |
40 |
| - bucket: str |
| 33 | +class StreamInfoBase: |
| 34 | + def __init__(self, file: str): |
| 35 | + self._file = file |
41 | 36 |
|
| 37 | + @property |
| 38 | + def file(self) -> str: |
| 39 | + return self._file |
42 | 40 |
|
43 |
| -@dataclass |
44 |
| -class StreamInfoPath: |
45 |
| - '''Contains information on accessing ServiceX data via a local Path |
46 |
| - ''' |
47 |
| - path: Path |
48 |
| - file: str |
49 | 41 |
|
| 42 | +class StreamInfoUrl(StreamInfoBase): |
| 43 | + def __init__(self, file: str, url: str, bucket: str): |
| 44 | + super().__init__(file) |
| 45 | + self._url = url |
| 46 | + self._bucket = bucket |
50 | 47 |
|
51 |
| -@dataclass |
52 |
| -class StreamInfoData: |
53 |
| - '''Contains information on accessing ServiceX data via converted data |
54 |
| - ''' |
55 |
| - data: Any |
56 |
| - file: str |
| 48 | + @property |
| 49 | + def url(self) -> str: |
| 50 | + return self._url |
| 51 | + |
| 52 | + @property |
| 53 | + def bucket(self) -> str: |
| 54 | + return self._bucket |
| 55 | + |
| 56 | + |
| 57 | +class StreamInfoPath(StreamInfoBase): |
| 58 | + def __init__(self, file: str, path: Path): |
| 59 | + super().__init__(file) |
| 60 | + self._path = path |
| 61 | + |
| 62 | + @property |
| 63 | + def path(self) -> Path: |
| 64 | + return self._path |
| 65 | + |
| 66 | + |
| 67 | +class StreamInfoData(StreamInfoBase): |
| 68 | + def __init__(self, file: str, data: Any): |
| 69 | + super().__init__(file) |
| 70 | + self._data = data |
| 71 | + |
| 72 | + @property |
| 73 | + def data(self) -> Any: |
| 74 | + return self._data |
57 | 75 |
|
58 | 76 |
|
59 | 77 | class ServiceXDataset(ServiceXABC):
|
@@ -320,7 +338,7 @@ async def _stream_url_buckets(self, selection_query: str, data_format: str) \
|
320 | 338 |
|
321 | 339 | # Reflect the files back up a level.
|
322 | 340 | async for r in minio_files:
|
323 |
| - yield StreamInfoUrl(minio_adaptor.get_access_url(request_id, r), r, request_id) |
| 341 | + yield StreamInfoUrl(r, minio_adaptor.get_access_url(request_id, r), request_id) |
324 | 342 |
|
325 | 343 | # Cache the final status
|
326 | 344 | await self._update_query_status(client, request_id)
|
@@ -397,7 +415,7 @@ async def _stream_return(self, selection_query: str,
|
397 | 415 | data Data converted to the "proper" format, depending
|
398 | 416 | on the converter call.
|
399 | 417 | '''
|
400 |
| - as_data = (StreamInfoData(await asyncio.ensure_future(converter(f.path)), f.file) |
| 418 | + as_data = (StreamInfoData(f.file, await asyncio.ensure_future(converter(f.path))) |
401 | 419 | async for f in self._stream_local_files(selection_query, data_format))
|
402 | 420 |
|
403 | 421 | async for r in as_data:
|
@@ -434,7 +452,7 @@ async def _stream_local_files(self, selection_query: str,
|
434 | 452 | self._get_files(selection_query, data_format, notifier))
|
435 | 453 |
|
436 | 454 | async for name, a_path in as_files:
|
437 |
| - yield StreamInfoPath(Path(await a_path), name) |
| 455 | + yield StreamInfoPath(name, Path(await a_path)) |
438 | 456 |
|
439 | 457 | async def _get_files(self, selection_query: str, data_type: str,
|
440 | 458 | notifier: _status_update_wrapper) \
|
|
0 commit comments