From d33457b7bc094aabeeaf4f1d4fe51e7adcf3b92c Mon Sep 17 00:00:00 2001 From: Jan Bouwhuis Date: Wed, 24 Nov 2021 15:15:27 +0100 Subject: [PATCH] Add bytes support for bitwise template operations (#60211) * Add bytes support for bitwise template operations * spelling * Update bitwise tests * remove try block for bytes conversion * do not accept empty `bytes` object --- homeassistant/helpers/template.py | 19 ++++++--- homeassistant/util/__init__.py | 25 ++++++++++++ tests/helpers/test_template.py | 66 +++++++++++++++++++++++++++++++ tests/util/test_init.py | 16 ++++++++ 4 files changed, 121 insertions(+), 5 deletions(-) diff --git a/homeassistant/helpers/template.py b/homeassistant/helpers/template.py index 07219942f98..3e2e3a55c3e 100644 --- a/homeassistant/helpers/template.py +++ b/homeassistant/helpers/template.py @@ -52,7 +52,12 @@ from homeassistant.helpers import ( ) from homeassistant.helpers.typing import TemplateVarsType from homeassistant.loader import bind_hass -from homeassistant.util import convert, dt as dt_util, location as loc_util +from homeassistant.util import ( + convert, + convert_to_int, + dt as dt_util, + location as loc_util, +) from homeassistant.util.async_ import run_callback_threadsafe from homeassistant.util.thread import ThreadWithException @@ -1629,14 +1634,18 @@ def regex_findall(value, find="", ignorecase=False): return re.findall(find, value, flags) -def bitwise_and(first_value, second_value): +def bitwise_and(first_value, second_value, little_endian=False): """Perform a bitwise and operation.""" - return first_value & second_value + return convert_to_int(first_value, little_endian=little_endian) & convert_to_int( + second_value, little_endian=little_endian + ) -def bitwise_or(first_value, second_value): +def bitwise_or(first_value, second_value, little_endian=False): """Perform a bitwise or operation.""" - return first_value | second_value + return convert_to_int(first_value, little_endian=little_endian) | convert_to_int( + second_value, little_endian=little_endian + ) def base64_encode(value): diff --git a/homeassistant/util/__init__.py b/homeassistant/util/__init__.py index 3840958d580..81f1389b47c 100644 --- a/homeassistant/util/__init__.py +++ b/homeassistant/util/__init__.py @@ -109,6 +109,31 @@ def convert( return default +def convert_to_int( + value: Any, default: int | None = None, little_endian: bool = False +) -> int | None: + """Convert value or bytes to int, returns default if fails. + + This supports bitwise integer operations on `bytes` objects. + By default the conversion is in Big-endian style (The last byte contains the least significant bit). + In Little-endian style the first byte contains the least significant bit. + """ + if isinstance(value, int): + return value + if isinstance(value, bytes) and value: + bytes_value = bytearray(value) + return_value = 0 + while len(bytes_value): + return_value <<= 8 + if little_endian: + return_value |= bytes_value.pop(len(bytes_value) - 1) + else: + return_value |= bytes_value.pop(0) + + return return_value + return convert(value, int, default=default) + + def ensure_unique_string( preferred_string: str, current_strings: Iterable[str] | KeysView[str] ) -> str: diff --git a/tests/helpers/test_template.py b/tests/helpers/test_template.py index d6649b058e2..a50884b71c8 100644 --- a/tests/helpers/test_template.py +++ b/tests/helpers/test_template.py @@ -1452,6 +1452,24 @@ def test_bitwise_and(hass): ) assert tpl.async_render() == 8 & 2 + tpl = template.Template( + """ +{{ ( value_a ) | bitwise_and(value_b) }} + """, + hass, + ) + variables = {"value_a": b"\x9b\xc2", "value_b": 0xFF00} + assert tpl.async_render(variables=variables) == 0x9B00 + + tpl = template.Template( + """ +{{ ( value_a ) | bitwise_and(value_b, little_endian=True) }} + """, + hass, + ) + variables = {"value_a": b"\xc2\x9b", "value_b": 0xFFFF} + assert tpl.async_render(variables=variables) == 0x9BC2 + def test_bitwise_or(hass): """Test bitwise_or method.""" @@ -1477,6 +1495,54 @@ def test_bitwise_or(hass): ) assert tpl.async_render() == 8 | 2 + tpl = template.Template( + """ +{{ value_a | bitwise_or(value_b) }} + """, + hass, + ) + variables = { + "value_a": b"\xc2\x9b", + "value_b": 0xFFFF, + } + assert tpl.async_render(variables=variables) == 65535 # 39874 + + tpl = template.Template( + """ +{{ ( value_a ) | bitwise_or(value_b) }} + """, + hass, + ) + variables = { + "value_a": 0xFF00, + "value_b": b"\xc2\x9b", + } + assert tpl.async_render(variables=variables) == 0xFF9B + + tpl = template.Template( + """ +{{ ( value_a ) | bitwise_or(value_b) }} + """, + hass, + ) + variables = { + "value_a": b"\xc2\x9b", + "value_b": 0x0000, + } + assert tpl.async_render(variables=variables) == 0xC29B + + tpl = template.Template( + """ +{{ ( value_a ) | bitwise_or(value_b, little_endian=True) }} + """, + hass, + ) + variables = { + "value_a": b"\xc2\x9b", + "value_b": 0, + } + assert tpl.async_render(variables=variables) == 0x9BC2 + def test_distance_function_with_1_state(hass): """Test distance function with 1 state.""" diff --git a/tests/util/test_init.py b/tests/util/test_init.py index 7a4f13cb767..0634f3cb6cf 100644 --- a/tests/util/test_init.py +++ b/tests/util/test_init.py @@ -98,6 +98,22 @@ def test_convert(): assert util.convert(object, int, 1) == 1 +def test_convert_to_int(): + """Test convert of bytes and numbers to int.""" + assert util.convert_to_int(b"\x9b\xc2") == 39874 + assert util.convert_to_int(b"") is None + assert util.convert_to_int(b"\x9b\xc2", 10) == 39874 + assert util.convert_to_int(b"\xc2\x9b", little_endian=True) == 39874 + assert util.convert_to_int(b"\xc2\x9b", 10, little_endian=True) == 39874 + assert util.convert_to_int("abc", 10) == 10 + assert util.convert_to_int("11.0", 10) == 10 + assert util.convert_to_int("12", 10) == 12 + assert util.convert_to_int("\xc2\x9b", 10) == 10 + assert util.convert_to_int(None, 10) == 10 + assert util.convert_to_int(None) is None + assert util.convert_to_int("NOT A NUMBER", 1) == 1 + + def test_ensure_unique_string(): """Test ensure_unique_string.""" assert util.ensure_unique_string("Beer", ["Beer", "Beer_2"]) == "Beer_3"