pw_status: Functions for checking the status code

- Add IsErrorCode() functions to simplify checking for specific status
  codes.
- Automatically replace some cases of status == Status::Code() with
  status.IsCode() in update_style.py.

Change-Id: I13be834a2413d694d41a7755d2e317c35d34c5bf
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/19521
Reviewed-by: Keir Mierle <keir@google.com>
Reviewed-by: Ewout van Bekkum <ewout@google.com>
Commit-Queue: Wyatt Hepler <hepler@google.com>

1 files changed, 16 insertions, 5 deletions

diff --git a/pw_status/py/pw_status/update_style.py b/pw_status/py/pw_status/update_style.py
index 57c92bbf0..bdff48363 100755
--- a/pw_status/py/pw_status/update_style.py
+++ b/pw_status/py/pw_status/update_style.py
@@ -45,6 +45,16 @@ _REMAP = {
     'DATA_LOSS': 'DataLoss',
 }
 
+_CODES = '|'.join(_REMAP.keys())
+_FUNCTIONS = '|'.join(_REMAP.values())
+
+_STATUS_WITH_SIZE_CTOR = re.compile(
+    fr'\bStatusWithSize\(Status::({_CODES}),\s*')
+_STATUS = re.compile(fr'\b(Status|StatusWithSize)::({_CODES})(?!")\b')
+_STATUS_EQUALITY = re.compile(
+    fr'Status::(?P<l_func>{_FUNCTIONS})\(\)\s+==\s+(?P<value>[a-zA-Z0-9_.()]+)|'
+    fr'\s+==\s+Status::(?P<r_func>{_FUNCTIONS})\(\)')
+
 
 def _remap_status_with_size(match) -> str:
     return f'StatusWithSize::{_REMAP[match.group(1)]}('
@@ -54,11 +64,10 @@ def _remap_codes(match) -> str:
     return f'{match.group(1)}::{_REMAP[match.group(2)]}()'
 
 
-_CODES = '|'.join(_REMAP.keys())
-
-_STATUS_WITH_SIZE_CTOR = re.compile(
-    fr'\bStatusWithSize\(Status::({_CODES}),\s*')
-_STATUS = re.compile(fr'\b(Status|StatusWithSize)::({_CODES})(?!")\b')
+def _remap_equality(match) -> str:
+    l_func, status, r_func = match.groups('')
+    func = l_func or r_func
+    return f'{status}.ok()' if func == 'Ok' else f'{status}.Is{func}()'
 
 
 def _parse_args():
@@ -92,6 +101,8 @@ def update_status(paths: Iterable[Path]) -> None:
             # Replace Status and StatusAWithSize
             text = _STATUS.sub(_remap_codes, text)
 
+            text = _STATUS_EQUALITY.sub(_remap_equality, text)
+
             if orig != text:
                 updated += 1
                 file.write_text(text)