From e2a39a7fa9fa25b30d0c5a0414f17245798e3d5d Mon Sep 17 00:00:00 2001
From: Ben Davies <ben.davies@anu.edu.au>
Date: Tue, 1 Oct 2024 15:04:27 +1000
Subject: [PATCH] Refactor set_item_codes() & remove old tests.

See issue #98 for reasoning.
---
 test/test_um2netcdf.py | 25 -------------------------
 umpost/um2netcdf.py    | 18 ++++++++++++------
 2 files changed, 12 insertions(+), 31 deletions(-)

diff --git a/test/test_um2netcdf.py b/test/test_um2netcdf.py
index 7aedce3..aa17e01 100644
--- a/test/test_um2netcdf.py
+++ b/test/test_um2netcdf.py
@@ -377,31 +377,6 @@ def test_stash_code_to_item_code_conversion():
     assert result == 30255
 
 
-def test_set_item_codes():
-    cube0 = DummyCube(1002, "d0", {um2nc.STASH: DummyStash(1, 2)})
-    cube1 = DummyCube(3004, "d1", {um2nc.STASH: DummyStash(3, 4)})
-    cubes = [cube0, cube1]
-
-    for cube in cubes:
-        assert hasattr(cube, um2nc.ITEM_CODE)
-
-    um2nc.set_item_codes(cubes)
-    c0, c1 = cubes
-
-    assert c0.item_code == 1002
-    assert c1.item_code == 3004
-
-
-def test_set_item_codes_avoid_overwrite():
-    item_code = 1007
-    item_code2 = 51006
-
-    cubes = [DummyCube(item_code, "fake_var"), DummyCube(item_code2, "fake_var2")]
-    um2nc.set_item_codes(cubes)
-    assert cubes[0].item_code == item_code
-    assert cubes[1].item_code == item_code2
-
-
 def add_stash(cube, stash):
     d = {um2nc.STASH: stash}
     setattr(cube, "attributes", d)
diff --git a/umpost/um2netcdf.py b/umpost/um2netcdf.py
index 147e98b..bde2beb 100644
--- a/umpost/um2netcdf.py
+++ b/umpost/um2netcdf.py
@@ -648,13 +648,19 @@ def to_stash_code(item_code: int):
 
 
 def set_item_codes(cubes):
-    for cube in cubes:
-        if hasattr(cube, ITEM_CODE):
-            msg = f"Cube {cube.var_name} already has 'item_code' attribute, skipping."
-            warnings.warn(msg)
-            continue
+    """
+    Add item code attribute to given cubes.
 
-        # hack: manually store item_code in cubes
+    Iris cube objects lack a item_code attribute, a single integer value
+    representing the combined stash/section code. This function converts the
+    cube's own stash/section code and stores as an "item_code" attribute. This
+    function is hacky from dynamically modifying the cube interface at runtime,
+    """
+    # TODO: should this be _set_item_codes() to flag as an internal detail?
+    for cube in cubes:
+        # NB: expanding the interface at runtime is somewhat hacky, however iris
+        # cube objects are defined in a 3rd party project. The alternative is
+        # passing primitives or additional data structures in process().
         item_code = to_item_code(cube.attributes[STASH])
         setattr(cube, ITEM_CODE, item_code)