From 0a92d7ace743f296ccb52699a35df492320cde25 Mon Sep 17 00:00:00 2001
From: Jacob Tomlinson <jacob@tomlinson.email>
Date: Tue, 1 Oct 2024 11:01:19 +0000
Subject: [PATCH] Add kubectl-ng cordon/uncordon

---
 .../kubectl-ng/kubectl_ng/_cordon_uncordon.py | 44 +++++++++++++++++++
 examples/kubectl-ng/kubectl_ng/cli.py         |  3 ++
 kr8s/_objects.py                              | 12 +++--
 3 files changed, 56 insertions(+), 3 deletions(-)
 create mode 100644 examples/kubectl-ng/kubectl_ng/_cordon_uncordon.py

diff --git a/examples/kubectl-ng/kubectl_ng/_cordon_uncordon.py b/examples/kubectl-ng/kubectl_ng/_cordon_uncordon.py
new file mode 100644
index 00000000..2a32f9ac
--- /dev/null
+++ b/examples/kubectl-ng/kubectl_ng/_cordon_uncordon.py
@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, Kr8s Developers (See LICENSE for list)
+# SPDX-License-Identifier: BSD 3-Clause License
+
+import typer
+from rich.console import Console
+
+from kr8s.asyncio.objects import Node
+
+console = Console()
+
+
+# Missing Options
+# TODO --dry-run='none'
+# TODO -l, --selector=''
+
+
+async def cordon(
+    node: str = typer.Argument(..., help="NODE"),
+):
+    """Mark node as unschedulable.
+
+    Examples:
+        # Mark node "foo" as unschedulable
+        kubectl-ng cordon foo
+    """
+    nodes = [await Node.get(node)]
+    for node_instance in nodes:
+        await node_instance.cordon()
+        console.print(f"node/{node_instance.name} cordoned")
+
+
+async def uncordon(
+    node: str = typer.Argument(..., help="NODE"),
+):
+    """Mark node as schedulable.
+
+    Examples:
+        # Mark node "foo" as schedulable
+        kubectl-ng uncordon foo
+    """
+    nodes = [await Node.get(node)]
+    for node_instance in nodes:
+        await node_instance.uncordon()
+        console.print(f"node/{node_instance.name} uncordoned")
diff --git a/examples/kubectl-ng/kubectl_ng/cli.py b/examples/kubectl-ng/kubectl_ng/cli.py
index 0751cae6..3342e243 100644
--- a/examples/kubectl-ng/kubectl_ng/cli.py
+++ b/examples/kubectl-ng/kubectl_ng/cli.py
@@ -5,6 +5,7 @@
 from ._api_resources import api_resources
 from ._api_versions import api_versions
 from ._config import config
+from ._cordon_uncordon import cordon, uncordon
 from ._create import create
 from ._delete import delete
 from ._exec import kexec
@@ -16,6 +17,8 @@
 app = typer.Typer(no_args_is_help=True)
 register(app, api_resources)
 register(app, api_versions)
+register(app, cordon)
+register(app, uncordon)
 register(app, create)
 register(app, delete)
 register(app, get)
diff --git a/kr8s/_objects.py b/kr8s/_objects.py
index cb7497f5..73dc2320 100644
--- a/kr8s/_objects.py
+++ b/kr8s/_objects.py
@@ -13,6 +13,7 @@
     BinaryIO,
     List,
     Literal,
+    Sequence,
     cast,
 )
 
@@ -21,6 +22,7 @@
 import jsonpath
 import yaml
 from box import Box
+from typing_extensions import Self
 
 import kr8s
 import kr8s.asyncio
@@ -239,7 +241,7 @@ async def get(
         field_selector: str | dict[str, str] | None = None,
         timeout: int = 2,
         **kwargs,
-    ) -> APIObject:
+    ) -> Self:
         """Get a Kubernetes resource by name or via selectors."""
         if api is None:
             if cls._asyncio:
@@ -283,6 +285,7 @@ async def get(
                 raise ValueError(
                     f"Expected exactly one {cls.kind} object. Use selectors to narrow down the search."
                 )
+            assert isinstance(resources[0], cls)
             return resources[0]
         raise NotFoundError(
             f"Could not find {cls.kind} {name} in namespace {namespace}."
@@ -685,7 +688,7 @@ def gen(cls, *args, **kwargs):
 
     # Must be the last method defined due to https://github.com/python/mypy/issues/17517
     @classmethod
-    async def list(cls, **kwargs) -> APIObject | list[APIObject]:
+    async def list(cls, **kwargs) -> Sequence[Self]:
         """List objects in Kubernetes.
 
         Args:
@@ -695,7 +698,10 @@ async def list(cls, **kwargs) -> APIObject | list[APIObject]:
             A list of objects.
         """
         api = await kr8s.asyncio.api()
-        return await api.async_get(kind=cls, **kwargs)
+        resources = await api.async_get(kind=cls, **kwargs)
+        if not isinstance(resources, list):
+            resources = [resources]
+        return [resource for resource in resources if isinstance(resource, cls)]
 
 
 ## v1 objects