From f708b6d159735faa1ce1d4671952bdd4b9971f4b Mon Sep 17 00:00:00 2001 From: Mose Mueller Date: Tue, 13 Feb 2024 11:17:43 +0100 Subject: [PATCH] fixing type hints --- pydase_service_base/database/influxdb_session.py | 8 +++++++- pydase_service_base/database/postgres_session.py | 7 ++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/pydase_service_base/database/influxdb_session.py b/pydase_service_base/database/influxdb_session.py index 21df4cc..4b74c4e 100644 --- a/pydase_service_base/database/influxdb_session.py +++ b/pydase_service_base/database/influxdb_session.py @@ -3,6 +3,12 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING, Any, NamedTuple +try: + from typing import Self # type: ignore +except ImportError: + from typing_extensions import Self + + from confz import FileSource from influxdb_client import ( Bucket, @@ -77,7 +83,7 @@ class InfluxDBSession: self._write_api: WriteApi self._buckets_api: BucketsApi - def __enter__(self) -> InfluxDBSession: + def __enter__(self) -> Self: self._client = InfluxDBClient(url=self.url, token=self.token, org=self.org) self._write_api = self._client.write_api(write_options=SYNCHRONOUS) # type: ignore return self diff --git a/pydase_service_base/database/postgres_session.py b/pydase_service_base/database/postgres_session.py index 9d3ab30..ed2e95b 100644 --- a/pydase_service_base/database/postgres_session.py +++ b/pydase_service_base/database/postgres_session.py @@ -6,6 +6,11 @@ import logging import re from typing import TYPE_CHECKING, Any +try: + from typing import Self # type: ignore +except ImportError: + from typing_extensions import Self + from confz import FileSource from dateutil.parser import ParserError, parse # type: ignore from sqlmodel import Session, SQLModel, create_engine @@ -155,7 +160,7 @@ class PostgresDatabaseSession(Session): ) ) - def __enter__(self) -> PostgresDatabaseSession: + def __enter__(self) -> Self: """Begins the runtime context related to the database session.""" super().__enter__()