Repository: canonical/operator
Pull request #1383 (2024-09-20 05:55:23)

fix: fix type of StatusBase subclasses by calling StatusBase.register in __init_subclass__
This is a simple alternative to #1382, following @dimaqq's suggestion to just not use `StatusBase.register` as a decorator. In this PR we instead do so via `__init_subclass__`.

james-garner-canonical: Thanks, @tonyandrewmeyer. I also tried `ag 'class .*\(.*StatusBase'` to see if anyone was subclassing `StatusBase` and didn't find anything.

If we wanted to deprecate register, we could move the functionality to `_register`, call that in `__init_subclass__` and `register`, and add a deprecation warning to `register`. If there's interest in doing that, I'm happy to add it to the PR.

Patch:
diff --git a/ops/model.py b/ops/model.py
index 707ef42..6405e70 100644
--- a/ops/model.py
+++ b/ops/model.py
@@ -1907,6 +1907,9 @@ class StatusBase:
             raise TypeError('cannot instantiate a base class')
         self.message = message
 
+    def __init_subclass__(cls):
+        StatusBase.register(cls)
+
     def __eq__(self, other: 'StatusBase') -> bool:
         if not isinstance(self, type(other)):
             return False
@@ -1965,7 +1968,6 @@ class StatusBase:
         return max(statuses, key=lambda status: cls._priorities.get(status.name, 0))
 
 
-@StatusBase.register
 class UnknownStatus(StatusBase):
     """The unit status is unknown.
 
@@ -1986,7 +1988,6 @@ class UnknownStatus(StatusBase):
         return 'UnknownStatus()'
 
 
-@StatusBase.register
 class ErrorStatus(StatusBase):
     """The unit status is error.
 
@@ -2000,7 +2001,6 @@ class ErrorStatus(StatusBase):
     name = 'error'
 
 
-@StatusBase.register
 class ActiveStatus(StatusBase):
     """The unit or application is ready and active.
 
@@ -2016,7 +2016,6 @@ class ActiveStatus(StatusBase):
         super().__init__(message)
 
 
-@StatusBase.register
 class BlockedStatus(StatusBase):
     """The unit or application requires manual intervention.
 
@@ -2027,7 +2026,6 @@ class BlockedStatus(StatusBase):
     name = 'blocked'
 
 
-@StatusBase.register
 class MaintenanceStatus(StatusBase):
     """The unit or application is performing maintenance tasks.
 
@@ -2041,7 +2039,6 @@ class MaintenanceStatus(StatusBase):
     name = 'maintenance'
 
 
-@StatusBase.register
 class WaitingStatus(StatusBase):
     """The unit or application is waiting on a charm it's integrated with.
 


Test patch:
diff --git a/test/test_model.py b/test/test_model.py
index 0eadd4d..7af3226 100644
--- a/test/test_model.py
+++ b/test/test_model.py
@@ -820,11 +820,15 @@ class TestModel:
         with pytest.raises(TypeError):
             ops.StatusBase('test')
 
-        class NoNameStatus(ops.StatusBase):
-            pass
+        with pytest.raises(TypeError):
+            # TypeError due to missing `name` attribute
+            class NoNameStatus(ops.StatusBase):  # pyright: ignore[reportUnusedClass]
+                pass
 
         with pytest.raises(TypeError):
-            ops.StatusBase.register(NoNameStatus)
+            # TypeError due to non str type `name` attribute
+            class NonStringNameStatus(ops.StatusBase):  # pyright: ignore[reportUnusedClass]
+                name = None  # pyright: ignore[reportAssignmentType]
 
     def test_status_repr(self):
         test_cases = {
diff --git a/test/test_testing.py b/test/test_testing.py
index 4c88f23..965ebc8 100644
--- a/test/test_testing.py
+++ b/test/test_testing.py
@@ -3232,6 +3232,9 @@ class TestHarness:
 
     def test_evaluate_status(self):
         class TestCharm(ops.CharmBase):
+            app_status_to_add: ops.StatusBase
+            unit_status_to_add: ops.StatusBase
+
             def __init__(self, framework: ops.Framework):
                 super().__init__(framework)
                 self.framework.observe(self.on.collect_app_status, self._on_collect_app_status)