From 2c6f3d40edce3eac96bcc33ec1ed1d2556284535 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Thu, 12 Oct 2023 08:12:39 -1000 Subject: [PATCH] Add optional basic cython implementation for frame_helper (#564) --- .github/workflows/ci.yml | 35 ++++++---- README.rst | 6 ++ aioesphomeapi/_frame_helper/base.pxd | 21 ++++++ aioesphomeapi/_frame_helper/base.py | 17 +++-- aioesphomeapi/_frame_helper/noise.pxd | 27 ++++++++ aioesphomeapi/_frame_helper/noise.py | 4 +- aioesphomeapi/_frame_helper/plain_text.pxd | 23 +++++++ aioesphomeapi/_frame_helper/plain_text.py | 9 +-- aioesphomeapi/connection.py | 4 +- pyproject.toml | 4 ++ setup.py | 76 +++++++++++++++++----- 11 files changed, 187 insertions(+), 39 deletions(-) create mode 100644 aioesphomeapi/_frame_helper/base.pxd create mode 100644 aioesphomeapi/_frame_helper/noise.pxd create mode 100644 aioesphomeapi/_frame_helper/plain_text.pxd diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6541399..e70a655 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -15,7 +15,7 @@ concurrency: jobs: ci: - name: ${{ matrix.name }} py ${{ matrix.python-version }} on ${{ matrix.os }} + name: ${{ matrix.name }} py ${{ matrix.python-version }} on ${{ matrix.os }} (${{ matrix.extension }}) runs-on: ${{ matrix.os }} strategy: fail-fast: false @@ -26,7 +26,10 @@ jobs: - "3.11" - "3.12" os: - - ubuntu-latest + - ubuntu-latest + extension: + - "skip_cython" + - "use_cython" steps: - uses: actions/checkout@v4 - name: Set up Python @@ -43,10 +46,20 @@ jobs: uses: actions/cache@v3 with: path: ${{ steps.pip-cache.outputs.dir }} - key: pip-${{ steps.python.outputs.python-version }}-${{ hashFiles('requirements.txt', 'requirements_test.txt') }} + key: pip-${{ steps.python.outputs.python-version }}-${{ matrix.extension }}-${{ hashFiles('requirements.txt', 'requirements_test.txt') }} restore-keys: | - pip-${{ steps.python.outputs.python-version }}- - - name: Set up Python environment + pip-${{ steps.python.outputs.python-version }}-${{ matrix.extension }}- + - name: Set up Python environment (no cython) + if: ${{ matrix.extension == 'skip_cython' }} + env: + SKIP_CYTHON: 1 + run: | + pip3 install -r requirements.txt -r requirements_test.txt + pip3 install -e . + - name: Set up Python environment (cython) + if: ${{ matrix.extension == 'use_cython' }} + env: + REQUIRE_CYTHON: 1 run: | pip3 install -r requirements.txt -r requirements_test.txt pip3 install -e . @@ -60,19 +73,19 @@ jobs: - run: flake8 aioesphomeapi name: Lint with flake8 - if: ${{ matrix.python-version == '3.11' }} + if: ${{ matrix.python-version == '3.11' && matrix.extension == 'skip_cython' }} - run: pylint aioesphomeapi name: Lint with pylint - if: ${{ matrix.python-version == '3.11' }} + if: ${{ matrix.python-version == '3.11' && matrix.extension == 'skip_cython' }} - run: black --check --diff --color aioesphomeapi tests name: Check formatting with black - if: ${{ matrix.python-version == '3.11' }} + if: ${{ matrix.python-version == '3.11' && matrix.extension == 'skip_cython' }} - run: isort --check --diff aioesphomeapi tests name: Check import order with isort - if: ${{ matrix.python-version == '3.11' }} + if: ${{ matrix.python-version == '3.11' && matrix.extension == 'skip_cython' }} - run: mypy aioesphomeapi name: Check typing with mypy - if: ${{ matrix.python-version == '3.11' }} + if: ${{ matrix.python-version == '3.11' && matrix.extension == 'skip_cython' }} - run: pytest -vv --tb=native tests name: Run tests with pytest - run: | @@ -86,4 +99,4 @@ jobs: exit 1 fi name: Check protobuf files match - if: ${{ matrix.python-version == '3.11' }} + if: ${{ matrix.python-version == '3.11' && matrix.extension == 'skip_cython' }} diff --git a/README.rst b/README.rst index e4c4003..1c92317 100644 --- a/README.rst +++ b/README.rst @@ -12,6 +12,12 @@ The module is available from the `Python Package Index bytearray | None: + def _read_exactly(self, length: _int) -> bytearray | None: """Read exactly length bytes from the buffer or None if all the bytes are not yet available.""" original_pos = self._pos new_pos = original_pos + length @@ -106,14 +108,15 @@ class APIFrameHelper(asyncio.Protocol): self._on_error(exc) def connection_lost(self, exc: Exception | None) -> None: + """Handle the connection being lost.""" self._handle_error( exc or SocketClosedAPIError(f"{self._log_name}: Connection lost") ) - return super().connection_lost(exc) def eof_received(self) -> bool | None: + """Handle EOF received.""" self._handle_error(SocketClosedAPIError(f"{self._log_name}: EOF received")) - return super().eof_received() + return False def close(self) -> None: """Close the connection.""" @@ -121,3 +124,9 @@ class APIFrameHelper(asyncio.Protocol): self._transport.close() self._transport = None self._writer = None + + def pause_writing(self) -> None: + """Stub.""" + + def resume_writing(self) -> None: + """Stub.""" diff --git a/aioesphomeapi/_frame_helper/noise.pxd b/aioesphomeapi/_frame_helper/noise.pxd new file mode 100644 index 0000000..17990b4 --- /dev/null +++ b/aioesphomeapi/_frame_helper/noise.pxd @@ -0,0 +1,27 @@ +import cython + +from .base cimport APIFrameHelper + + +cdef object TYPE_CHECKING + +cdef class APINoiseFrameHelper(APIFrameHelper): + + cdef object _noise_psk + cdef object _expected_name + cdef object _state + cdef object _dispatch + cdef object _server_name + cdef object _proto + cdef object _decrypt + cdef object _encrypt + cdef bint _is_ready + + @cython.locals( + header=bytearray, + preamble=cython.uint, + msg_size_high=cython.uint, + msg_size_low=cython.uint, + end_of_frame_pos=cython.uint, + ) + cpdef data_received(self, bytes data) \ No newline at end of file diff --git a/aioesphomeapi/_frame_helper/noise.py b/aioesphomeapi/_frame_helper/noise.py index 79863a8..c9a56bd 100644 --- a/aioesphomeapi/_frame_helper/noise.py +++ b/aioesphomeapi/_frame_helper/noise.py @@ -144,7 +144,9 @@ class APINoiseFrameHelper(APIFrameHelper): header = self._read_exactly(3) if header is None: return - preamble, msg_size_high, msg_size_low = header + preamble = header[0] + msg_size_high = header[1] + msg_size_low = header[2] if preamble != 0x01: self._handle_error_and_close( ProtocolAPIError( diff --git a/aioesphomeapi/_frame_helper/plain_text.pxd b/aioesphomeapi/_frame_helper/plain_text.pxd new file mode 100644 index 0000000..656a27f --- /dev/null +++ b/aioesphomeapi/_frame_helper/plain_text.pxd @@ -0,0 +1,23 @@ +import cython + +from .base cimport APIFrameHelper + + +cdef object TYPE_CHECKING +cdef object WRITE_EXCEPTIONS +cdef object bytes_to_varuint, varuint_to_bytes + +cdef class APIPlaintextFrameHelper(APIFrameHelper): + + @cython.locals( + msg_type=bytes, + length=bytes, + init_bytes=bytearray, + add_length=bytearray, + end_of_frame_pos=cython.uint, + length_int=cython.uint, + preamble=cython.uint, + length_high=cython.uint, + maybe_msg_type=cython.uint + ) + cpdef data_received(self, bytes data) \ No newline at end of file diff --git a/aioesphomeapi/_frame_helper/plain_text.py b/aioesphomeapi/_frame_helper/plain_text.py index dcd6f0f..059bca8 100644 --- a/aioesphomeapi/_frame_helper/plain_text.py +++ b/aioesphomeapi/_frame_helper/plain_text.py @@ -50,8 +50,10 @@ class APIPlaintextFrameHelper(APIFrameHelper): if init_bytes is None: return msg_type_int: int | None = None - length_int: int | None = None - preamble, length_high, maybe_msg_type = init_bytes + length_int = 0 + preamble = init_bytes[0] + length_high = init_bytes[1] + maybe_msg_type = init_bytes[2] if preamble != 0x00: if preamble == 0x01: self._handle_error_and_close( @@ -88,7 +90,7 @@ class APIPlaintextFrameHelper(APIFrameHelper): if add_length is None: return length += add_length - length_int = bytes_to_varuint(length) + length_int = bytes_to_varuint(length) or 0 # Since the length is longer than 1 byte we do not have the # message type yet. msg_type = b"" @@ -105,7 +107,6 @@ class APIPlaintextFrameHelper(APIFrameHelper): msg_type_int = bytes_to_varuint(msg_type) if TYPE_CHECKING: - assert length_int is not None assert msg_type_int is not None if length_int == 0: diff --git a/aioesphomeapi/connection.py b/aioesphomeapi/connection.py index 395d6fa..f1bcf90 100644 --- a/aioesphomeapi/connection.py +++ b/aioesphomeapi/connection.py @@ -325,7 +325,7 @@ class APIConnection: assert self._socket is not None if self._params.noise_psk is None: - _, fh = await loop.create_connection( + _, fh = await loop.create_connection( # type: ignore[type-var] lambda: APIPlaintextFrameHelper( on_pkt=process_packet, on_error=self._report_fatal_error, @@ -337,7 +337,7 @@ class APIConnection: else: noise_psk = self._params.noise_psk assert noise_psk is not None - _, fh = await loop.create_connection( + _, fh = await loop.create_connection( # type: ignore[type-var] lambda: APINoiseFrameHelper( noise_psk=noise_psk, expected_name=self._params.expected_name, diff --git a/pyproject.toml b/pyproject.toml index 532f8bc..fe9b15c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,3 +29,7 @@ disable = [ "duplicate-code", "too-many-lines", ] + +[build-system] +requires = ['setuptools>=65.4.1', 'wheel', 'Cython>=3.0.2'] + diff --git a/setup.py b/setup.py index f5ed665..a62ae0a 100644 --- a/setup.py +++ b/setup.py @@ -3,6 +3,9 @@ import os from setuptools import find_packages, setup +import os +from distutils.command.build_ext import build_ext + here = os.path.abspath(os.path.dirname(__file__)) @@ -31,20 +34,59 @@ DOWNLOAD_URL = "{}/archive/{}.zip".format(GITHUB_URL, VERSION) with open(os.path.join(here, "requirements.txt")) as requirements_txt: REQUIRES = requirements_txt.read().splitlines() -setup( - name=PROJECT_PACKAGE_NAME, - version=VERSION, - url=PROJECT_URL, - download_url=DOWNLOAD_URL, - author=PROJECT_AUTHOR, - author_email=PROJECT_EMAIL, - description="Python API for interacting with ESPHome devices.", - long_description=long_description, - license=PROJECT_LICENSE, - packages=find_packages(exclude=["tests", "tests.*"]), - include_package_data=True, - zip_safe=False, - install_requires=REQUIRES, - python_requires=">=3.9", - test_suite="tests", -) + +setup_kwargs = { + "name": PROJECT_PACKAGE_NAME, + "version": VERSION, + "url": PROJECT_URL, + "download_url": DOWNLOAD_URL, + "author": PROJECT_AUTHOR, + "author_email": PROJECT_EMAIL, + "description": "Python API for interacting with ESPHome devices.", + "long_description": long_description, + "license": PROJECT_LICENSE, + "packages": find_packages(exclude=["tests", "tests.*"]), + "include_package_data": True, + "zip_safe": False, + "install_requires": REQUIRES, + "python_requires": ">=3.9", + "test_suite": "tests", +} + + +class OptionalBuildExt(build_ext): + def build_extensions(self): + try: + super().build_extensions() + except Exception: + pass + + +def cythonize_if_available(setup_kwargs): + if os.environ.get("SKIP_CYTHON", False): + return + try: + from Cython.Build import cythonize + + setup_kwargs.update( + dict( + ext_modules=cythonize( + [ + "aioesphomeapi/_frame_helper/plain_text.py", + "aioesphomeapi/_frame_helper/noise.py", + "aioesphomeapi/_frame_helper/base.py", + ], + compiler_directives={"language_level": "3"}, # Python 3 + ), + cmdclass=dict(build_ext=OptionalBuildExt), + ) + ) + except Exception: + if os.environ.get("REQUIRE_CYTHON"): + raise + pass + + +cythonize_if_available(setup_kwargs) + +setup(**setup_kwargs)