making feature extra bits more generic (#13)

* making feature extra bits more generic * adds a features testcase for phoenix invoices * use Union for < py3.9 support
lnbits · Jul 18, 2023 · 4854335 · 4854335
1 parent 87ad682
commit 4854335
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 48 deletions.
diff --git a/bolt11/models/features.py b/bolt11/models/features.py
@@ -1,6 +1,6 @@
 from enum import Enum
 from math import floor
-from typing import Dict, NamedTuple, Optional
+from typing import Dict, NamedTuple, Optional, Union
 
 from bitstring import BitArray, Bits
 
@@ -29,61 +29,42 @@ class Feature(Enum):
     option_scid_alias = 16
     option_payment_metadata = 17
     option_zeroconf_chanids = 18
-    extra_1 = 19
-    extra_2 = 20
-    extra_3 = 21
-    extra_4 = 22
-    extra_5 = 23
-    extra_6 = 24
-    extra_7 = 25
-    extra_8 = 26
-    extra_9 = 27
-    extra_10 = 28
-    extra_11 = 29
-    extra_12 = 30
-    extra_13 = 31
-    extra_14 = 32
-    extra_15 = 33
-    extra_16 = 34
-    extra_17 = 35
-    extra_18 = 36
-    extra_19 = 37
-    extra_20 = 38
-    extra_21 = 39
-    extra_22 = 40
-    extra_23 = 41
-    extra_24 = 42
-    extra_25 = 43
-    extra_26 = 44
-    extra_27 = 45
-    extra_28 = 46
-    extra_29 = 47
-    extra_30 = 48
-    extra_31 = 49
+
+
+class FeatureExtra:
+    def __init__(self, index: int):
+        self.feature_index = index
+
+    @property
+    def value(self) -> int:
+        return self.feature_index + len(Feature)
+
+    @property
+    def name(self) -> str:
+        return f"extra_{self.feature_index - len(Feature)}"
 
 
 class Features(NamedTuple):
     data: Bits
-    feature_list: Dict[Feature, FeatureState]
+    feature_list: Dict[Union[Feature, FeatureExtra], FeatureState]
 
     @classmethod
     def from_bitstring(cls, data: Bits) -> "Features":
         length = data.length
-        feature_list: Dict[Feature, FeatureState] = {}
+        feature_list: Dict[Union[Feature, FeatureExtra], FeatureState] = {}
         for i in range(0, length):
             feature_index = floor(i / 2)
-            if floor(i / 2) > len(Feature):
-                raise ValueError(f"Feature index ({i}) out of range, word_length: {length}")
             si = i + 1
             cut = data[-si : -si + 1] if i > 0 else data[-si:]
             if bool(cut):
-                feature = Feature(floor(i / 2))
-                if feature not in feature_list:
-                    feature_list[Feature(feature_index)] = FeatureState.supported if i % 2 else FeatureState.required
+                feature: Union[Feature, FeatureExtra] = (
+                    Feature(feature_index) if feature_index < len(Feature) else FeatureExtra(feature_index)
+                )
+                feature_list[feature] = FeatureState.supported if i % 2 else FeatureState.required
         return cls(data, feature_list)
 
     @classmethod
-    def from_feature_list(cls, feature_list: Dict[Feature, FeatureState]) -> "Features":
+    def from_feature_list(cls, feature_list: Dict[Union[Feature, FeatureExtra], FeatureState]) -> "Features":
         length = max([feature.value + 1 for feature in feature_list]) * 2
         data = BitArray(length=length)
         for feature, feature_state in feature_list.items():

diff --git a/tests/test_bolt11_examples.py b/tests/test_bolt11_examples.py
@@ -1,7 +1,7 @@
 from bolt11.decode import decode
 from bolt11.encode import encode
 from bolt11.models.fallback import Fallback
-from bolt11.models.features import Feature, Features, FeatureState
+from bolt11.models.features import Feature, FeatureExtra, Features, FeatureState
 from bolt11.models.routehint import RouteHint
 from bolt11.types import Bolt11
 
@@ -726,7 +726,7 @@ def test_example_11(self):
             "feature_list": {
                 Feature.var_onion_optin: FeatureState.required,
                 Feature.payment_secret: FeatureState.required,
-                Feature.extra_31: FeatureState.supported,
+                FeatureExtra(31): FeatureState.supported,
             },
             "signature": (
                 "5755469bf4b8e6b6ae7a1308d5f9bad5c82812e0855cd24fac242aa323fa820c5c551ede4faeabcb7fb6d5a46"
@@ -789,7 +789,7 @@ def test_example_12(self):
             "feature_list": {
                 Feature.var_onion_optin: FeatureState.required,
                 Feature.payment_secret: FeatureState.required,
-                Feature.extra_31: FeatureState.supported,
+                FeatureExtra(31): FeatureState.supported,
             },
             "signature": (
                 "5755469bf4b8e6b6ae7a1308d5f9bad5c82812e0855cd24fac242aa323fa820c5c551ede4faeabcb7fb6d5a46"
@@ -859,7 +859,7 @@ def test_example_13(self):
             "feature_list": {
                 Feature.var_onion_optin: FeatureState.required,
                 Feature.payment_secret: FeatureState.required,
-                Feature.extra_31: FeatureState.supported,
+                FeatureExtra(31): FeatureState.supported,
             },
             "signature": (
                 "150a5252308f25bc2641a186de87470189bb003774326beee33b9a2a720d1584386631c5dda6fc3"
@@ -905,7 +905,7 @@ def test_example_14(self):
             "feature_list": {
                 Feature.var_onion_optin: FeatureState.required,
                 Feature.payment_secret: FeatureState.required,
-                Feature.extra_6: FeatureState.required,
+                FeatureExtra(6): FeatureState.required,
             },
             "signature": (
                 "f5d27be7d9c27d3aa521bc35d77cabd6bda18f1f61716445b19e27e4e17a887508ea8de5a8e1d94f561248f65"

diff --git a/tests/test_features.py b/tests/test_features.py
@@ -1,7 +1,7 @@
 import pytest
 
 from bolt11.decode import decode
-from bolt11.models.features import Feature, Features, FeatureState
+from bolt11.models.features import Feature, FeatureExtra, Features, FeatureState
 
 
 class TestDecodeFeatures:
@@ -33,7 +33,7 @@ class TestDecodeFeatures:
                 {
                     Feature.var_onion_optin: FeatureState.required,
                     Feature.payment_secret: FeatureState.required,
-                    Feature.extra_31: FeatureState.supported,
+                    FeatureExtra(31): FeatureState.supported,
                 },
             ),
             (
@@ -47,7 +47,29 @@ class TestDecodeFeatures:
                 {
                     Feature.var_onion_optin: FeatureState.required,
                     Feature.payment_secret: FeatureState.required,
-                    Feature.extra_6: FeatureState.required,
+                    FeatureExtra(6): FeatureState.required,
+                },
+            ),
+            (
+                # phoenix invoice
+                (
+                    "lnbc1950n1pjtrgxnpp5ye3lhh8ye8ywm85evshn9wyhdsdg9a350tdhm3dyw89mwfht8s9qdqqxqyj"
+                    "w5q9q7sqqqqqqqqqqqqqqqqqqqqqqqqq9qsqsp5z22wjrrm0lgl32e0yes38dzmvjxnajrvanhw3hp4"
+                    "duq4k55wl0gsrzjqwryaup9lh50kkranzgcdnn2fgvx390wgj5jd07rwr3vxeje0glclludsryzx7vv"
+                    "vqqqqqlgqqqqqeqqjqqq3zq0d9kw8q4fhsgxh595f2l0ass4zaj2pdknzhxzzrlf7g5wgsk3nlgzzed"
+                    "uhnp6mva9jehwcq9y4hrllwt6ffl822q5drdgvxtjspmgfnls"
+                ),
+                {
+                    "var_onion_optin": "supported",
+                    "payment_secret": "supported",
+                    "basic_mpp": "supported",
+                    "extra_56": "supported",
+                },
+                {
+                    Feature.var_onion_optin: FeatureState.supported,
+                    Feature.payment_secret: FeatureState.supported,
+                    Feature.basic_mpp: FeatureState.supported,
+                    FeatureExtra(56): FeatureState.supported,
                 },
             ),
         ],