From eb99b6fdca089a14e1b5867dc4d1a0988191fbba Mon Sep 17 00:00:00 2001 From: Dom Sekotill Date: Wed, 18 May 2022 01:50:05 +0100 Subject: [PATCH] Make Container.started return the same type as self --- behave_utils/docker.py | 2 +- behave_utils/mysql.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/behave_utils/docker.py b/behave_utils/docker.py index 7f40768..0d682ae 100644 --- a/behave_utils/docker.py +++ b/behave_utils/docker.py @@ -222,7 +222,7 @@ class Container(Item): self.stop(rm=True) @contextmanager - def started(self) -> Iterator[Container]: + def started(self: T) -> Iterator[T]: """ Return a context manager that ensures the container is started when the context is entered """ diff --git a/behave_utils/mysql.py b/behave_utils/mysql.py index fdf8ca5..13c98b7 100644 --- a/behave_utils/mysql.py +++ b/behave_utils/mysql.py @@ -16,6 +16,7 @@ from time import sleep from typing import TYPE_CHECKING from typing import Iterator from typing import Sequence +from typing import TypeVar from behave import fixture @@ -39,6 +40,8 @@ class Mysql(Container): Container subclass for a database container """ + T = TypeVar('T', bound='Mysql') + def __init__( self, version: str = "latest", @@ -95,7 +98,7 @@ class Mysql(Container): return Cli(self, "mysqldump") @contextmanager - def started(self) -> Iterator[Container]: + def started(self: T) -> Iterator[T]: """ Return a context manager that only enters once the database is initialised """ -- GitLab