Add a neuro storage du command

CHANGELOG.D/2201.feature

@@ -0,0 +1 @@
+Added `neuro storage du` command that allows to retrieve cluster's storage disk usage.

CLI.md

@@ -40,6 +40,7 @@
 		* [neuro storage mkdir](#neuro-storage-mkdir)
 		* [neuro storage mv](#neuro-storage-mv)
 		* [neuro storage tree](#neuro-storage-tree)
+		* [neuro storage du](#neuro-storage-du)
 	* [neuro image](#neuro-image)
 		* [neuro image ls](#neuro-image-ls)
 		* [neuro image push](#neuro-image-push)
@@ -1011,6 +1012,7 @@ Name | Description|
 | _[neuro storage mkdir](#neuro-storage-mkdir)_| Make directories |
 | _[neuro storage mv](#neuro-storage-mv)_| Move or rename files and directories |
 | _[neuro storage tree](#neuro-storage-tree)_| List contents of directories in a tree-like format |
+| _[neuro storage du](#neuro-storage-du)_| Show current usage of storage |
 
 
 
@@ -1239,6 +1241,25 @@ Name | Description|
 
 
 
+### neuro storage du
+
+Show current usage of storage.
+
+**Usage:**
+
+```bash
+neuro storage du [OPTIONS]
+```
+
+**Options:**
+
+Name | Description|
+|----|------------|
+|_--help_|Show this message and exit.|
+
+
+
+
 ## neuro image
 
 Container image operations.

neuro-cli/docs/storage.md

@@ -20,6 +20,7 @@ Storage operations.
 | [_mkdir_](storage.md#mkdir) | Make directories |
 | [_mv_](storage.md#mv) | Move or rename files and directories |
 | [_tree_](storage.md#tree) | List contents of directories in a... |
+| [_du_](storage.md#du) | Show current usage of storage |
 
 
 ### cp
@@ -303,3 +304,24 @@ home dir (storage:)
 | _--sort \[name | size | time\]_ | sort by given field, default is name |
 
 
+
+### du
+
+Show current usage of storage
+
+
+#### Usage
+
+```bash
+neuro storage du [OPTIONS]
+```
+
+Show current usage of storage.
+
+#### Options
+
+| Name | Description |
+| :--- | :--- |
+| _--help_ | Show this message and exit. |
+
+

neuro-cli/src/neuro_cli/formatters/storage.py

@@ -42,6 +42,7 @@
     StorageProgressStart,
     StorageProgressStep,
 )
+from neuro_sdk.storage import DiskUsageInfo
 from neuro_sdk.url_utils import _extract_path
 
 from neuro_cli.root import Root
@@ -852,3 +853,12 @@ def _human_readable(self, size: int) -> str:
 
     def _none(self, size: int) -> str:
         return ""
+
+
+class DiskUsageFormatter:
+    def __call__(self, usage: DiskUsageInfo) -> Text:
+        ret = Text.from_markup(f"Disk usage for cluster [b]{usage.cluster_name}[/]:\n")
+        ret.append(f"Total: {format_size(usage.total)}\n")
+        ret.append(f"Used:  {format_size(usage.used)}\n")
+        ret.append(f"Free:  {format_size(usage.free)}")
+        return ret

neuro-cli/src/neuro_cli/storage.py

@@ -17,6 +17,7 @@
 from .formatters.storage import (
     BaseFilesFormatter,
     DeleteProgress,
+    DiskUsageFormatter,
     FilesSorter,
     LongFilesFormatter,
     SimpleFilesFormatter,
@@ -203,6 +204,15 @@ async def glob(root: Root, patterns: Sequence[str]) -> None:
                 root.print(file)
 
 
+@command()
+async def du(root: Root) -> None:
+    """
+    Show current usage of storage.
+    """
+    usage = await root.client.storage.disk_usage()
+    root.print(DiskUsageFormatter()(usage))
+
+
 class FileFilterParserOption(click.parser.Option):
     def process(self, value: str, state: click.parser.ParsingState) -> None:
         assert self.action == "append"
@@ -722,6 +732,7 @@ async def _is_dir(root: Root, uri: URL) -> bool:
 storage.add_command(mkdir)
 storage.add_command(mv)
 storage.add_command(tree)
+storage.add_command(du)
 
 
 async def calc_filters(

neuro-cli/tests/unit/formatters/ascii/TestUsageFormatter.test_formatter_0.ref

@@ -0,0 +1,4 @@
+Disk usage for cluster default:
+Total: 97.7K
+Used:  78.1K
+Free:  19.5K

neuro-cli/tests/unit/formatters/test_storage_formatters.py

@@ -5,11 +5,13 @@
 from yarl import URL
 
 from neuro_sdk import Action, FileStatus, FileStatusType
+from neuro_sdk.storage import DiskUsageInfo
 
 from neuro_cli.formatters.storage import (
     BaseFilesFormatter,
     BSDAttributes,
     BSDPainter,
+    DiskUsageFormatter,
     FilesSorter,
     GnuIndicators,
     GnuPainter,
@@ -342,3 +344,12 @@ def test_sorter(self) -> None:
             self.files[1],
             self.files[2],
         ]
+
+
+class TestUsageFormatter:
+    def test_formatter(self, rich_cmp: Any) -> None:
+        usage = DiskUsageInfo(
+            total=100000, used=80000, free=20000, cluster_name="default"
+        )
+        formatter = DiskUsageFormatter()
+        rich_cmp(formatter(usage))

neuro-sdk/docs/storage_reference.rst

@@ -122,6 +122,14 @@ Storage
 
       :return: data structure for given *uri*, :class:`FileStatus` object.
 
+   .. comethod:: disk_usage(cluster_name: Optional[str] = None) -> DiskUsageInfo
+
+      Return information about disk usage in given cluster.
+
+      :param str cluster_name: cluster name to retrieve info. If ``None`` current
+                               cluster will be used.
+
+      :return: data structure for given cluster, :class:`DiskUsageInfo` object.
 
    .. rubric:: File operations
 
@@ -526,3 +534,28 @@ StorageProgress event classes
    .. attribute:: is_dir
 
       **True** if removed item was a directory; **False** otherwise. :class:`bool`
+
+
+DiskUsageInfo
+=============
+
+.. class:: DiskUsageInfo
+
+   *Read-only* :class:`~dataclasses.dataclass` for describing disk usage in particular
+   cluster
+
+   .. attribute:: cluster_name
+
+      Name of cluster, :class:`str`.
+
+   .. attribute:: total
+
+      Total storage size in bytes, :class:`int`.
+
+   .. attribute:: used
+
+      Used storage size in bytes, :class:`int`.
+
+   .. attribute:: free
+
+      Free storage size in bytes, :class:`int`.

neuro-sdk/src/neuro_sdk/storage.py

@@ -98,6 +98,14 @@ def name(self) -> str:
         return Path(self.path).name
 
 
+@dataclass(frozen=True)
+class DiskUsageInfo:
+    cluster_name: str
+    total: int
+    used: int
+    free: int
+
+
 class Storage(metaclass=NoPublicConstructor):
     def __init__(self, core: _Core, config: Config) -> None:
         self._core = core
@@ -331,6 +339,22 @@ async def stat(self, uri: URL) -> FileStatus:
             res = await resp.json()
             return _file_status_from_api_stat(uri.host, res["FileStatus"])
 
+    async def disk_usage(self, cluster_name: Optional[str] = None) -> DiskUsageInfo:
+        cluster_name = cluster_name or self._config.cluster_name
+        uri = self._normalize_uri(
+            URL(f"storage://{cluster_name}/{self._config.username}")
+        )
+        assert uri.host is not None
+        url = self._get_storage_url(uri, normalized=True)
+        url = url.with_query(op="GETDISKUSAGE")
+        auth = await self._config._api_auth()
+
+        request_time = time.time()
+        async with self._core.request("GET", url, auth=auth) as resp:
+            self._set_time_diff(request_time, resp)
+            res = await resp.json()
+            return _disk_usage_from_api(uri.host, res)
+
     @asyncgeneratorcontextmanager
     async def open(
         self, uri: URL, offset: int = 0, size: Optional[int] = None
@@ -912,6 +936,15 @@ def _file_status_from_api_stat(cluster_name: str, values: Dict[str, Any]) -> Fil
     )
 
 
+def _disk_usage_from_api(cluster_name: str, values: Dict[str, Any]) -> DiskUsageInfo:
+    return DiskUsageInfo(
+        cluster_name=cluster_name,
+        total=values["total"],
+        used=values["used"],
+        free=values["free"],
+    )
+
+
 def _parse_content_range(rng_str: Optional[str]) -> slice:
     if rng_str is None:
         raise RuntimeError("Missed header Content-Range")

neuro-sdk/tests/test_storage.py

@@ -24,7 +24,7 @@
     StorageProgressStep,
 )
 from neuro_sdk.abc import StorageProgressDelete
-from neuro_sdk.storage import _parse_content_range
+from neuro_sdk.storage import DiskUsageInfo, _parse_content_range
 
 from tests import _RawTestServerFactory, _TestServerFactory
 
@@ -289,6 +289,46 @@ async def handler(request: web.Request) -> web.StreamResponse:
     ]
 
 
+async def test_storage_disk_usage(
+    aiohttp_server: _TestServerFactory, make_client: _MakeClient
+) -> None:
+    async def handler(request: web.Request) -> web.StreamResponse:
+        assert "b3" in request.headers
+        assert request.path == "/storage/user"
+        assert request.query == {"op": "GETDISKUSAGE"}
+        return web.json_response({"total": 100, "used": 20, "free": 80})
+
+    app = web.Application()
+    app.router.add_get("/storage/user", handler)
+
+    srv = await aiohttp_server(app)
+
+    async with make_client(srv.make_url("/")) as client:
+        res = await client.storage.disk_usage()
+
+    assert res == DiskUsageInfo(total=100, used=20, free=80, cluster_name="default")
+
+
+async def test_storage_disk_usage_another_cluster(
+    aiohttp_server: _TestServerFactory, make_client: _MakeClient
+) -> None:
+    async def handler(request: web.Request) -> web.StreamResponse:
+        assert "b3" in request.headers
+        assert request.path == "/storage2/user"
+        assert request.query == {"op": "GETDISKUSAGE"}
+        return web.json_response({"total": 100, "used": 20, "free": 80})
+
+    app = web.Application()
+    app.router.add_get("/storage2/user", handler)
+
+    srv = await aiohttp_server(app)
+
+    async with make_client(srv.make_url("/")) as client:
+        res = await client.storage.disk_usage(cluster_name="another")
+
+    assert res == DiskUsageInfo(total=100, used=20, free=80, cluster_name="another")
+
+
 async def test_storage_ls_another_cluster(
     aiohttp_server: _TestServerFactory, make_client: _MakeClient
 ) -> None: