robshakir · JoseIgnacioTamayo · Sep 28, 2023 · Nov 23, 2021
diff --git a/docs/yang.md b/docs/yang.md
@@ -53,9 +53,9 @@ PyangBind does not currently try and be feature complete against the YANG langua
  **Type**            | **Sub-Statement**   | **Supported Type**      | **Unit Tests**
  --------------------|--------------------|--------------------------|---------------
  **binary**          | -                   | [bytes](https://docs.python.org/3/library/stdtypes.html?#bytes)           | tests/binary
- \-                   | length              | Supported           | tests/binary
- **bits**            | -                   | Not supported           | N/A
- \-                   | position            | Not supported           | N/A
+ \-                  | length              | Supported           | tests/binary
+ **bits**            | -                   | Set\[str\]              | tests/bits
+ \-                  | position            | Supported               |
  **boolean**         | -                   | YANGBool                | tests/boolean-empty
  **empty**           | -                   | YANGBool                | tests/boolean-empty
  **decimal64**       | -                   | [Decimal](https://docs.python.org/2/library/decimal.html) | tests/decimal64

diff --git a/pyangbind/lib/serialise.py b/pyangbind/lib/serialise.py
@@ -200,6 +200,8 @@ def map_pyangbind_type(self, map_val, original_yang_type, obj):
             return self.preprocess_element(obj.get())
         elif map_val in ["decimal.Decimal"]:
             return self.yangt_decimal(obj)
+        elif map_val in ["YANGBits"]:
+            return six.text_type(obj)
 
     def yangt_int(self, obj):
         # for values that are 32-bits and under..
@@ -686,7 +688,12 @@ def load_json(
                 else:
                     # use the set method
                     pass
-            elif pybind_attr in ["RestrictedClassType", "ReferencePathType", "RestrictedPrecisionDecimal"]:
+            elif pybind_attr in [
+                "RestrictedClassType",
+                "ReferencePathType",
+                "RestrictedPrecisionDecimal",
+                "YANGBits",
+            ]:
                 # normal but valid types - which use the std set method
                 pass
             elif pybind_attr is None:

diff --git a/pyangbind/lib/yangtypes.py b/pyangbind/lib/yangtypes.py
@@ -1370,3 +1370,76 @@ def __repr__(self):
 
     def __str__(self, encoding="ascii", errors="replace"):
         return str(self, encoding=encoding, errors=errors)
+
+
+def YANGBitsType(allowed_bits):
+    class YANGBits(set):
+        """Map the ``bits`` built-in type of YANG
+
+        From RFC 6020, 9.7:
+
+        > The bits built-in type represents a bit set. That is, a bits value
+        > is a set of flags identified by small integer position numbers
+        > starting at 0. Each bit number has an assigned name.
+        >
+        > ... In the canonical form, the bit values are separated by a single
+        > space character and they appear ordered by their position.
+
+        __init__ parses such a string, and __str__() prints the value as
+        above. In the Python model, the bits that are set are kept as a set
+        of strings.
+
+        Override the ``add()``, ``discard()`` and similar methods to set and
+        clear the bits, checking for legal values.
+
+        :param allowed_bits: dictionary of legal bit values and positions; it
+            is set when generating the YANG-Python type mapping, when
+            ``pyangbing.lib.pybing.build_elemtype`` recurses over the YANG model
+        """
+
+        _pybind_generated_by = "YANGBits"
+
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+            self._allowed_bits = allowed_bits
+
+            if args:
+                value = args[0]
+                for bit in value.split():
+                    self.add(bit)
+
+        def _add_bit_definition(self, bit, position):
+            self._allowed_bits[bit] = position
+
+        # overwrite set methods to 1/ check for legal values and 2/ set the
+        # changed flag
+        def add(self, bit):
+            if bit not in self._allowed_bits:
+                raise ValueError(f"Bit value {bit} not valid, expected one of {self._allowed_bits}")
+            super().add(bit)
+            self._mchanged = True
+
+        def clear(self):
+            super().clear()
+            self._mchanged = True
+
+        def discard(self, bit):
+            if bit not in self._allowed_bits:
+                raise ValueError(f"Bit value {bit} not valid, expected one of {self._allowed_bits}")
+            super().discard(bit)
+            self._mchanged = True
+
+        def pop(self):
+            super().pop()
+            self._mchanged = True
+
+        def remove(self, bit):
+            super().remove(bit)
+            self._mchanged = True
+
+        def __str__(self, encoding="ascii", errors="replace"):
+            """Return bits as shown in JSON."""
+            sort_key = self._allowed_bits.__getitem__
+            return " ".join(sorted(self, key=sort_key))
+
+    return YANGBits
diff --git a/pyangbind/plugin/pybind.py b/pyangbind/plugin/pybind.py
@@ -34,7 +34,13 @@
 
 import pyangbind.helpers.misc as misc_help
 from pyangbind.helpers.identity import IdentityStore
-from pyangbind.lib.yangtypes import RestrictedClassType, YANGBool, safe_name, YANGBinary
+from pyangbind.lib.yangtypes import (
+    RestrictedClassType,
+    YANGBool,
+    safe_name,
+    YANGBinary,
+    YANGBitsType,
+)
 
 # Python3 support
 if six.PY3:
@@ -105,6 +111,12 @@
         "base_type": True,
         "pytype": RestrictedClassType(base_type=int, restriction_dict={"range": ["0..255"]}, int_size=8),
     },
+    "bits": {
+        "native_type": "YANGBitsType(allowed_bits={})",
+        "base_type": True,
+        "quote_arg": True,
+        "pytype": YANGBitsType(allowed_bits={}),
+    },
     "uint16": {
         "native_type": (
             "RestrictedClassType(base_type=int," + " restriction_dict={'range': ['0..65535']}," + "int_size=16)"
@@ -317,6 +329,7 @@ def build_pybind(ctx, modules, fd):
         "YANGDynClass",
         "ReferenceType",
         "YANGBinary",
+        "YANGBitsType",
     ]
     for library in yangtypes_imports:
         ctx.pybind_common_hdr += "from pyangbind.lib.yangtypes import {}\n".format(library)
@@ -496,6 +509,7 @@ def build_typedefs(ctx, defnd):
     known_types = list(class_map.keys())
     known_types.append("enumeration")
     known_types.append("leafref")
+    known_types.append("bits")
     base_types = copy.deepcopy(known_types)
     process_typedefs_ordered = []
 
@@ -575,6 +589,7 @@ def build_typedefs(ctx, defnd):
         # in the class_map, and hence we append it here.
         known_types.append("enumeration")
         known_types.append("leafref")
+        known_types.append("bits")
 
         # Don't allow duplicate definitions of types
         if type_name in known_types:
@@ -1327,6 +1342,22 @@ def build_elemtype(ctx, et, prefix=False):
                     pp.pprint(et.arg)
                     pp.pprint(base_stmt.arg)
                 sys.exit(127)
+        elif et.arg == "bits":
+            allowed_bits = {}
+            for bit in et.search("bit"):
+                position = bit.search_one("position")
+                if position is not None:
+                    pos = position.arg
+                else:
+                    pos = 1 + max(allowed_bits)
+                allowed_bits[bit.arg] = pos
+            cls = "restricted-bits"
+            elemtype = {
+                "native_type": f"YANGBitsType(allowed_bits={allowed_bits})",
+                "base_type": True,
+                "parent_type": "bits",
+                "quote_arg": True,
+            }
         else:
             # For all other cases, then we should be able to look up directly in the
             # class_map for the defined type, since these are not 'derived' types

diff --git a/tests/bits/__init__.py b/tests/bits/__init__.py
diff --git a/tests/bits/bits.yang b/tests/bits/bits.yang
@@ -0,0 +1,49 @@
+module bits {
+  yang-version 1.1;
+  namespace "http://rob.sh/yang/test/bits";
+  prefix "foo";
+  organization "BugReports Inc";
+  contact "A bug reporter";
+
+  description
+      "A test module";
+  revision 2014-01-01 {
+      description "april-fools";
+      reference "fooled-you";
+  }
+
+  // a typedef for bits, which is then applied to a leaf
+  typedef typedefed-bits {
+    type bits {
+      bit foo {
+	position 0;
+      }
+      bit bar {
+	position 1;
+      }
+      bit baz {
+	position 2;
+      }
+    }
+  }
+
+  leaf bits2 {
+    type typedefed-bits;
+  }
+
+  // a leaf containing the bits type definition
+  leaf mybits {
+    type bits {
+      bit flag1 {
+	position 0;
+      }
+      bit flag2 {
+	position 1;
+      }
+      bit flag3 {
+	position 2;
+      }
+    }
+    default "flag1 flag3";
+  }
+}
diff --git a/tests/bits/run.py b/tests/bits/run.py
@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+from __future__ import unicode_literals
+
+import unittest
+
+from tests.base import PyangBindTestCase
+
+
+class BitsTests(PyangBindTestCase):
+    yang_files = ["bits.yang"]
+
+    def setUp(self):
+        self.instance = self.bindings.bits()
+
+    def test_default_bits(self):
+        self.assertEqual(
+            self.instance.mybits._default,
+            set(["flag1", "flag3"]),
+        )
+        self.assertFalse(self.instance.mybits._changed())
+
+    def test_default_bits_with_in(self):
+        self.assertTrue(
+            "flag1" in self.instance.mybits._default
+            and "flag2" not in self.instance.mybits._default
+                and "flag3" in self.instance.mybits._default
+        )
+
+    def test_default_bits_is_unchanged(self):
+        self.assertFalse(self.instance.mybits._changed())
+
+    def test_set_flags(self):
+        for value, valid in [("flag1", True), ("flag2", True), ("unknownflag", False)]:
+            with self.subTest(value=value, valid=valid):
+                allowed = True
+                try:
+                    self.instance.mybits.add(value)
+                except ValueError:
+                    allowed = False
+                self.assertEqual(allowed, valid)
+
+    def test_check_flags_unset(self):
+        self.instance.mybits.discard("flag2")
+        self.assertTrue("flag2" not in self.instance.mybits)
+        self.assertTrue(self.instance.mybits._changed())
+
+    def test_check_flags_set(self):
+        self.instance.bits2.add("foo")
+        self.assertTrue("foo" in self.instance.bits2)
+        self.assertEqual(set(["foo"]), self.instance.bits2)
+        self.assertTrue(self.instance.bits2._changed())
+
+    def test_bits_position(self):
+        self.instance.bits2.add("bar")
+        self.instance.bits2.add("baz")
+        self.instance.bits2.add("foo")
+        self.assertEqual(str(self.instance.bits2), "foo bar baz")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/serialise/json-deserialise/json-deserialise.yang b/tests/serialise/json-deserialise/json-deserialise.yang
@@ -171,6 +171,20 @@ module json-deserialise {
           }
         }
       }
+
+      leaf bits {
+        type bits {
+	  bit flag1 {
+	    position 1;
+	  }
+	  bit flag2 {
+	    position 2;
+	  }
+	  bit flag3 {
+	    position 3;
+	  }
+	}
+      }
     }
 
     list t1 {
@@ -189,5 +203,4 @@ module json-deserialise {
       type string;
     }
   }
-
 }
diff --git a/tests/serialise/json-deserialise/json/alltypes.json b/tests/serialise/json-deserialise/json/alltypes.json
@@ -35,7 +35,8 @@
                 "int8": 1, 
                 "uint64": 1, 
                 "int64": 1, 
-                "restricted-string": "aardvark"
+                "restricted-string": "aardvark",
+		"bits": "flag1 flag2"
             }
         }
     }

diff --git a/tests/serialise/json-deserialise/run.py b/tests/serialise/json-deserialise/run.py
@@ -79,6 +79,7 @@ def test_all_the_types(self):
                         "uint64": 1,
                         "int64": 1,
                         "restricted-string": "aardvark",
+                        "bits": set(["flag1", "flag2"]),
                     }
                 },
                 "t1": {"32": {"target": "32"}, "16": {"target": "16"}},

diff --git a/tests/serialise/json-serialise/json-serialise.yang b/tests/serialise/json-serialise/json-serialise.yang
@@ -199,6 +199,20 @@ module json-serialise {
         type decimalrangetype;
       }
 
+      leaf bits {
+        type bits {
+	  bit flag1 {
+	    position 1;
+	  }
+	  bit flag2 {
+	    position 2;
+	  }
+	  bit flag3 {
+	    position 3;
+	  }
+	}
+      }
+
       choice test-choice {
         case one {
           leaf one-leaf {

diff --git a/tests/serialise/json-serialise/json/expected-output.json b/tests/serialise/json-serialise/json/expected-output.json
@@ -78,7 +78,8 @@
                 "decleaf": 42.4422,
                 "typedef-decimal": 21.21,
                 "range-decimal": 4.44443322,
-                "typedef-decimalrange": 42.42
+                "typedef-decimalrange": 42.42,
+                "bits": "flag1 flag3"
             }
         }
     }

diff --git a/tests/serialise/json-serialise/run.py b/tests/serialise/json-serialise/run.py
@@ -63,6 +63,8 @@ def test_full_serialise(self):
         self.serialise_obj.c1.l1[1].range_decimal = Decimal("4.44443322")
         self.serialise_obj.c1.l1[1].typedef_decimalrange = Decimal("42.42")
         self.serialise_obj.c1.l1[1].decleaf = Decimal("42.4422")
+        self.serialise_obj.c1.l1[1].bits.add("flag1")
+        self.serialise_obj.c1.l1[1].bits.add("flag3")
 
         for i in range(1, 10):
             self.serialise_obj.c1.l2.add(i)