waldiez 0.1.17__py3-none-any.whl → 0.1.18__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of waldiez might be problematic. Click here for more details.

waldiez/__init__.py CHANGED
@@ -1,17 +1,50 @@
1
1
  """Waldiez package."""
2
2
 
3
+ import logging
4
+
3
5
  from ._version import __version__
4
6
  from .conflict_checker import check_conflicts
5
7
  from .exporter import WaldiezExporter
6
8
  from .models import Waldiez
7
9
  from .runner import WaldiezRunner
8
10
 
11
+
12
+ # pylint: disable=too-few-public-methods
13
+ class FlamlFilter(logging.Filter):
14
+ """Filter out flaml.automl is not available message."""
15
+
16
+ def filter(self, record: logging.LogRecord) -> bool:
17
+ """Filter out flaml.automl is not available message.
18
+
19
+ this is just annoying:
20
+
21
+ ```
22
+ flaml.automl is not available.
23
+ Please install flaml[automl] to enable AutoML functionalities.
24
+ ```
25
+
26
+ Parameters
27
+ ----------
28
+ record : logging.LogRecord
29
+ Log record to filter.
30
+
31
+ Returns
32
+ -------
33
+ bool
34
+ Whether to filter out the log record.
35
+ """
36
+ return "flaml.automl is not available" not in record.getMessage()
37
+
38
+
9
39
  # flag to check if ag2 and autogen-agentchat
10
40
  # are installed at the same time
11
41
  __WALDIEZ_CHECKED_FOR_CONFLICTS = False
42
+ # flag to handle flaml logging
43
+ # suppress the annoying message about flaml.automl
44
+ __WALDIEZ_HANDLED_FLAML_LOGGING = False
12
45
 
13
46
 
14
- def _check_conflicts_once() -> None:
47
+ def _check_conflicts() -> None:
15
48
  """Check for conflicts once."""
16
49
  # pylint: disable=global-statement
17
50
  global __WALDIEZ_CHECKED_FOR_CONFLICTS
@@ -20,8 +53,18 @@ def _check_conflicts_once() -> None:
20
53
  __WALDIEZ_CHECKED_FOR_CONFLICTS = True
21
54
 
22
55
 
23
- _check_conflicts_once()
56
+ def _handle_flaml_logging() -> None:
57
+ """Handle flaml logging once."""
58
+ # pylint: disable=global-statement
59
+ global __WALDIEZ_HANDLED_FLAML_LOGGING
60
+ if __WALDIEZ_HANDLED_FLAML_LOGGING is False:
61
+ __WALDIEZ_HANDLED_FLAML_LOGGING = True
62
+ flam_logger = logging.getLogger("flaml")
63
+ flam_logger.addFilter(FlamlFilter())
64
+
24
65
 
66
+ _check_conflicts()
67
+ _handle_flaml_logging()
25
68
 
26
69
  __all__ = [
27
70
  "Waldiez",
waldiez/_version.py CHANGED
@@ -1,3 +1,3 @@
1
1
  """Version information for Waldiez."""
2
2
 
3
- __version__ = "0.1.17"
3
+ __version__ = "0.1.18"
waldiez/cli.py CHANGED
@@ -90,7 +90,7 @@ def run(
90
90
  raise typer.Exit(code=1) from error
91
91
  waldiez = Waldiez.from_dict(data)
92
92
  runner = WaldiezRunner(waldiez)
93
- results = runner.run(stream=None, output_path=output_path)
93
+ results = runner.run(output_path=output_path)
94
94
  logger = _get_logger()
95
95
  if isinstance(results, list):
96
96
  logger.info("Results:")
waldiez/runner.py CHANGED
@@ -12,11 +12,11 @@ import importlib.util
12
12
  import io
13
13
  import os
14
14
  import shutil
15
+ import site
15
16
  import subprocess # nosemgrep # nosec
16
17
  import sys
17
18
  import tempfile
18
19
  from contextlib import contextmanager
19
- from contextvars import ContextVar
20
20
  from pathlib import Path
21
21
  from types import TracebackType
22
22
  from typing import (
@@ -36,8 +36,6 @@ from .models.waldiez import Waldiez
36
36
  if TYPE_CHECKING:
37
37
  from autogen import ChatResult # type: ignore
38
38
 
39
- from .io import WaldiezIOStream
40
-
41
39
 
42
40
  @contextmanager
43
41
  def _chdir(to: Union[str, Path]) -> Iterator[None]:
@@ -61,6 +59,11 @@ def _chdir(to: Union[str, Path]) -> Iterator[None]:
61
59
  os.chdir(old_cwd)
62
60
 
63
61
 
62
+ def refresh_site_packages() -> None:
63
+ """Refresh the site packages."""
64
+ site.main()
65
+
66
+
64
67
  class WaldiezRunner:
65
68
  """Waldiez runner class."""
66
69
 
@@ -71,9 +74,6 @@ class WaldiezRunner:
71
74
  self._waldiez = waldiez
72
75
  self._running = False
73
76
  self._file_path = file_path
74
- self._stream: ContextVar[Optional["WaldiezIOStream"]] = ContextVar(
75
- "waldiez_stream", default=None
76
- )
77
77
  self._exporter = WaldiezExporter(waldiez)
78
78
 
79
79
  @classmethod
@@ -136,10 +136,6 @@ class WaldiezRunner:
136
136
  """Exit the context manager."""
137
137
  if self._running:
138
138
  self._running = False
139
- token = self._stream.get()
140
- if token is not None:
141
- self._stream.reset(token)
142
- del token
143
139
 
144
140
  @property
145
141
  def waldiez(self) -> Waldiez:
@@ -151,24 +147,13 @@ class WaldiezRunner:
151
147
  """Get the running status."""
152
148
  return self._running
153
149
 
154
- def _get_print_function(self) -> Callable[..., None]:
155
- """Get the print function."""
156
- token = self._stream.get()
157
- if token is not None:
158
- return token.print
159
- return print
160
-
161
- def _install_requirements(self) -> None:
150
+ def _install_requirements(self, printer: Callable[..., None]) -> None:
162
151
  """Install the requirements for the flow."""
163
152
  extra_requirements = set(
164
153
  req for req in self.waldiez.requirements if req not in sys.modules
165
154
  )
166
155
  if extra_requirements:
167
- print_function = self._get_print_function()
168
- # pylint: disable=inconsistent-quotes
169
- print_function(
170
- f"Installing requirements: {', '.join(extra_requirements)}"
171
- )
156
+ printer(f"Installing requirements: {', '.join(extra_requirements)}")
172
157
  pip_install = [sys.executable, "-m", "pip", "install"]
173
158
  if not in_virtualenv():
174
159
  pip_install.append("--user")
@@ -180,18 +165,23 @@ class WaldiezRunner:
180
165
  ) as proc:
181
166
  if proc.stdout:
182
167
  for line in io.TextIOWrapper(proc.stdout, encoding="utf-8"):
183
- print_function(line.strip())
168
+ printer(line.strip())
184
169
  if proc.stderr:
185
170
  for line in io.TextIOWrapper(proc.stderr, encoding="utf-8"):
186
- print_function(line.strip())
187
- print_function(
171
+ printer(line.strip())
172
+ printer("Refreshing site packages...")
173
+ refresh_site_packages()
174
+ printer(
188
175
  "Requirements installed.\n"
189
176
  "NOTE: If new packages were added and you are using Jupyter, "
190
177
  "you might need to restart the kernel."
191
178
  )
192
179
 
180
+ @staticmethod
193
181
  def _after_run(
194
- self, temp_dir: Path, output_path: Optional[Union[str, Path]]
182
+ temp_dir: Path,
183
+ output_path: Optional[Union[str, Path]],
184
+ printer: Callable[..., None],
195
185
  ) -> None:
196
186
  if output_path:
197
187
  destination_dir = Path(output_path).parent
@@ -202,9 +192,7 @@ class WaldiezRunner:
202
192
  )
203
193
  destination_dir.mkdir(parents=True, exist_ok=True)
204
194
  # copy the contents of the temp dir to the destination dir
205
- self._get_print_function()(
206
- f"Copying the results to {destination_dir}"
207
- )
195
+ printer(f"Copying the results to {destination_dir}")
208
196
  for item in temp_dir.iterdir():
209
197
  # skip cache files
210
198
  if (
@@ -257,6 +245,11 @@ class WaldiezRunner:
257
245
  Union[ChatResult, List[ChatResult]]
258
246
  The result(s) of the chat(s).
259
247
  """
248
+ # pylint: disable=import-outside-toplevel
249
+ from autogen.io import IOStream # type: ignore
250
+
251
+ printer = IOStream.get_default().print
252
+ self._install_requirements(printer)
260
253
  results: Union["ChatResult", List["ChatResult"]] = []
261
254
  if not uploads_root:
262
255
  uploads_root = Path(tempfile.mkdtemp())
@@ -284,32 +277,15 @@ class WaldiezRunner:
284
277
  old_vars = self._set_env_vars()
285
278
  module = importlib.util.module_from_spec(spec)
286
279
  spec.loader.exec_module(module)
287
- print_function = self._get_print_function()
288
- print_function("Starting workflow...")
280
+ printer("Starting workflow...")
289
281
  results = module.main()
290
282
  sys.path.pop(0)
291
283
  self._reset_env_vars(old_vars)
292
- self._after_run(temp_dir, output_path)
284
+ self._after_run(temp_dir, output_path, printer)
293
285
  return results
294
286
 
295
- def _run(
296
- self,
297
- output_path: Optional[Union[str, Path]],
298
- uploads_root: Optional[Union[str, Path]],
299
- ) -> Union["ChatResult", List["ChatResult"]]:
300
- self._install_requirements()
301
- token = self._stream.get()
302
- if token is not None:
303
- # pylint: disable=import-outside-toplevel
304
- from .io import WaldiezIOStream
305
-
306
- with WaldiezIOStream.set_default(token):
307
- return self._do_run(output_path, uploads_root)
308
- return self._do_run(output_path, uploads_root)
309
-
310
287
  def run(
311
288
  self,
312
- stream: Optional["WaldiezIOStream"] = None,
313
289
  output_path: Optional[Union[str, Path]] = None,
314
290
  uploads_root: Optional[Union[str, Path]] = None,
315
291
  ) -> Union["ChatResult", List["ChatResult"]]:
@@ -317,8 +293,6 @@ class WaldiezRunner:
317
293
 
318
294
  Parameters
319
295
  ----------
320
- stream : Optional[WaldiezIOStream], optional
321
- The stream to use, by default None.
322
296
  output_path : Optional[Union[str, Path]], optional
323
297
  The output path, by default None.
324
298
  uploads_root : Optional[Union[str, Path]], optional
@@ -337,14 +311,11 @@ class WaldiezRunner:
337
311
  if self._running is True:
338
312
  raise RuntimeError("Workflow already running")
339
313
  self._running = True
340
- token = self._stream.set(stream)
341
314
  file_path = output_path or self._file_path
342
315
  try:
343
- return self._run(file_path, uploads_root)
316
+ return self._do_run(file_path, uploads_root)
344
317
  finally:
345
318
  self._running = False
346
- self._stream.reset(token)
347
- del token
348
319
 
349
320
 
350
321
  def in_virtualenv() -> bool:
@@ -1,10 +1,11 @@
1
- Metadata-Version: 2.3
1
+ Metadata-Version: 2.4
2
2
  Name: waldiez
3
- Version: 0.1.17
3
+ Version: 0.1.18
4
4
  Summary: waldiez
5
5
  Project-URL: homepage, https://waldiez.github.io/waldiez/
6
6
  Project-URL: repository, https://github.com/waldiez/waldiez.git
7
7
  Author-email: Panagiotis Kasnesis <pkasnesis@thingenious.io>, Lazaros Toumanidis <laztoum@protonmail.com>, Stella Ioannidou <stella@humancentered.gr>
8
+ License-File: LICENSE
8
9
  Classifier: Development Status :: 3 - Alpha
9
10
  Classifier: Intended Audience :: Developers
10
11
  Classifier: Intended Audience :: Science/Research
@@ -60,7 +61,7 @@ Requires-Dist: mdx-include==1.4.2; extra == 'docs'
60
61
  Requires-Dist: mdx-truly-sane-lists==1.3; extra == 'docs'
61
62
  Requires-Dist: mkdocs-jupyter==0.25.1; extra == 'docs'
62
63
  Requires-Dist: mkdocs-macros-plugin==1.3.7; extra == 'docs'
63
- Requires-Dist: mkdocs-material==9.5.48; extra == 'docs'
64
+ Requires-Dist: mkdocs-material==9.5.49; extra == 'docs'
64
65
  Requires-Dist: mkdocs-minify-html-plugin==0.2.3; extra == 'docs'
65
66
  Requires-Dist: mkdocs==1.6.1; extra == 'docs'
66
67
  Requires-Dist: mkdocstrings-python==1.12.2; extra == 'docs'
@@ -88,7 +89,7 @@ To a python script or a jupyter notebook with the corresponding [ag2](https://gi
88
89
 
89
90
  - Convert .waldiez flows to .py or .ipynb
90
91
  - Run a .waldiez flow
91
- - Provide a custom [IOSStream](https://ag2ai.github.io/ag2/docs/reference/io/base#iostream) to handle input and output.
92
+ - Store the runtime logs of a flow to csv for further analysis
92
93
 
93
94
  ## Installation
94
95
 
@@ -178,73 +179,6 @@ runner = WaldiezRunner.load(flow_path)
178
179
  runner.run(output_path=output_path)
179
180
  ```
180
181
 
181
- #### Run a flow with a custom IOStream
182
-
183
- ```python
184
- # Run the flow with a custom IOStream
185
- # In case the standard 'input' and 'print' functions cannot be used
186
- import time
187
- import threading
188
-
189
- from typing import Any
190
-
191
- from waldiez import WaldiezRunner
192
- from waldiez.io import WaldiezIOStream
193
-
194
- flow_path = "/path/to/a/flow.waldiez"
195
- output_path = "/path/to/an/output.py"
196
-
197
-
198
- def custom_print_function(*args: Any, sep: str = " ", **kwargs: Any) -> None:
199
- """Custom print function."""
200
- print(*args, sep=sep, **kwargs)
201
-
202
-
203
- # Custom input handler
204
- class InputProcessorWrapper:
205
- """Wrapper input processor.
206
-
207
- To manage the interaction between the custom input processor and IOStream.
208
- """
209
-
210
- def __init__(self):
211
- self.stream = None # Placeholder for the WaldiezIOStream instance
212
- self.lock = threading.Lock() # Ensure thread-safe operations
213
-
214
- def custom_input_processor(self, prompt: str) -> None:
215
- """Simulate external input and send it back to the IOStream."""
216
- def external_input_simulation():
217
- with self.lock: # Ensure thread-safe access
218
- time.sleep(2) # Simulate delay for network input
219
- if self.stream:
220
- self.stream.set_input("Simulated external input")
221
- else:
222
- raise RuntimeError("Stream reference not set!")
223
-
224
- threading.Thread(target=external_input_simulation, daemon=True).start()
225
-
226
- def set_stream(self, stream: "WaldiezIOStream"):
227
- """Set the WaldiezIOStream instance."""
228
- with self.lock: # Ensure thread-safe setting of the stream
229
- self.stream = stream
230
-
231
- processor_wrapper = InputProcessorWrapper()
232
-
233
- stream = WaldiezIOStream(
234
- input_timeout=30,
235
- print_function=
236
- on_prompt_input=processor_wrapper.custom_input_processor,
237
- )
238
-
239
- # Link the processor wrapper to the WaldiezIOStream instance
240
- processor_wrapper.set_stream(custom_stream)
241
-
242
- with WaldiezIOStream.set_default(io_stream):
243
- runner = WaldiezRunner.load(flow_path)
244
- runner.run(stream=io_stream, output_path=output_path)
245
-
246
- ```
247
-
248
182
  ### Tools
249
183
 
250
184
  - [ag2 (formerly AutoGen)](https://github.com/ag2ai/ag2)
@@ -1,11 +1,11 @@
1
- waldiez/__init__.py,sha256=ASGmxP3JcASIv89_tMYhYxvW0FlHK1jM0Nfb1hMN-z0,722
1
+ waldiez/__init__.py,sha256=9BLMLUURndkGEwnSNF1RXg1XKm4WTYM3cUeuYo-M0qY,1936
2
2
  waldiez/__main__.py,sha256=mUQWu4CF03Jbbscvcfb_9EFTnIMliJJJmAuWf0sRRZU,110
3
- waldiez/_version.py,sha256=JZiuUXn7Efmws1BQsY7xqwPgrly1krkrpN-yh6tNz8s,63
4
- waldiez/cli.py,sha256=Ayp6Xu2HSBBFIouYh8cArjOjgXv5kbtZd_PYewcDklU,6347
3
+ waldiez/_version.py,sha256=sywGJ0P0JJRsNByFmXUw9vI_Pr9KX-kYoUgjhkhJKD8,63
4
+ waldiez/cli.py,sha256=19LFqep8urzcLm6TzVSyJEOKI2Q8XXZbuuFsgV0k9LQ,6334
5
5
  waldiez/conflict_checker.py,sha256=E-w0TfTivDAVpNvjd_NSBeaaFsWtWAyxKoSz0-g6x2U,880
6
6
  waldiez/exporter.py,sha256=iKe-l_Me8NRWsXHIdBcrOrnLT9XIyp4iYi4HLuuj2jA,9342
7
7
  waldiez/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
8
- waldiez/runner.py,sha256=v4A4NHdSnh96TQ5Vx1vs8UrTsshc2oAwDZrEHmq4RkU,11537
8
+ waldiez/runner.py,sha256=cFhiJI3LJKU0003DeyQaQcpvq6mZzJSWn9v5Fo4WxB4,10407
9
9
  waldiez/exporting/__init__.py,sha256=GMY7qTRpNmc7tpaCFNGLT5wX9eu26NvoNPuYX8MzP50,344
10
10
  waldiez/exporting/agents/__init__.py,sha256=v5KA112W_EFYwXE2TSBKYyO8rRKUOUAOpFS5CMSnfRs,110
11
11
  waldiez/exporting/agents/agent.py,sha256=d2ut2B90-aG4uBRhPWS7B4MKIvSJqQva6DjlXiqleVQ,7513
@@ -39,7 +39,6 @@ waldiez/exporting/utils/method_utils.py,sha256=7-RUMTylNM2W0iM1bPX2_Gn3553XZSl2s
39
39
  waldiez/exporting/utils/naming.py,sha256=VdoVODQduhXIs9hQFWUVEVqTaSyNDt7rkECsuIgXYwI,3196
40
40
  waldiez/exporting/utils/object_string.py,sha256=2kdIu4in3iUV92a2KbLWwp9POhvY-fzF_r2AGVnCKls,2166
41
41
  waldiez/exporting/utils/path_check.py,sha256=aO49sbk6hUHEr65UjNNUSKO5WCGnQitiT733W-kGVtI,1062
42
- waldiez/io/__init__.py,sha256=-w6ywIjG8ByxGtpSeRt1bjhnP4CZzpHf8WRWckjs4ng,5172
43
42
  waldiez/models/__init__.py,sha256=IMq8vzuAgmv92kHSSuZQLF38vVd31ojgOHonuHqWYj0,2888
44
43
  waldiez/models/waldiez.py,sha256=yI38WEXPJa-FmRF6Try-MabWoPCeqlxq8VUt0sD-fJg,9428
45
44
  waldiez/models/agents/__init__.py,sha256=3ZyVYBHMFzZjRMIdPrBF6HLg82LPAlEubL6utL6KhfU,1856
@@ -85,8 +84,8 @@ waldiez/models/model/model_data.py,sha256=VLPb60rJeZEgVZCjjkQGiwTrKz7OddVLrXCrEu
85
84
  waldiez/models/skill/__init__.py,sha256=rU88bajKOGMYoHFcE8MP0jW9H0MswbQmvz5wxS35BYE,169
86
85
  waldiez/models/skill/skill.py,sha256=fhsAI413an2_d4DBIkf7dzEuWk6rGs2t4sl97a4dj20,3473
87
86
  waldiez/models/skill/skill_data.py,sha256=RTWn8Od6w7g-nRIpsS29sqZ8sPm5dCPiK7-qXmU-KD4,815
88
- waldiez-0.1.17.dist-info/METADATA,sha256=xfStnh_3pvXf9OU87UM3VMr3l8mchce7g2-OnfCI3mg,10262
89
- waldiez-0.1.17.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
90
- waldiez-0.1.17.dist-info/entry_points.txt,sha256=9MQ8Y1rD19CU7UwjNPwoyTRpQsPs2QimjrtwTD0bD6k,44
91
- waldiez-0.1.17.dist-info/licenses/LICENSE,sha256=VQEHM6WMQLRu1qaGl3GWsoOknDwro-69eGo4NLIJPIM,1064
92
- waldiez-0.1.17.dist-info/RECORD,,
87
+ waldiez-0.1.18.dist-info/METADATA,sha256=8KXupJpMO4yQ62SbTcBJx4IrueJQx1RmOR_6yDQPwuE,8161
88
+ waldiez-0.1.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
89
+ waldiez-0.1.18.dist-info/entry_points.txt,sha256=9MQ8Y1rD19CU7UwjNPwoyTRpQsPs2QimjrtwTD0bD6k,44
90
+ waldiez-0.1.18.dist-info/licenses/LICENSE,sha256=VQEHM6WMQLRu1qaGl3GWsoOknDwro-69eGo4NLIJPIM,1064
91
+ waldiez-0.1.18.dist-info/RECORD,,
@@ -1,4 +1,4 @@
1
1
  Wheel-Version: 1.0
2
- Generator: hatchling 1.26.3
2
+ Generator: hatchling 1.27.0
3
3
  Root-Is-Purelib: true
4
4
  Tag: py3-none-any
waldiez/io/__init__.py DELETED
@@ -1,142 +0,0 @@
1
- """Custom IOStream class to use with autogen.
2
-
3
- It is meant to be used when we want to use custom
4
- `print` and `input`. For example, when a websocket
5
- is used to trigger a UI element that requires user input.
6
- and sends back the user's input to the websocket. In the same
7
- way, we can use it to forward what is meant to be printed.
8
- """
9
-
10
- import threading
11
- from typing import Any, Callable, Optional
12
-
13
- from autogen.io import IOStream # type: ignore[import-untyped]
14
-
15
-
16
- class WaldiezIOStream(IOStream):
17
- """Custom IOStream class to handle `print` and `input` functions."""
18
-
19
- def __init__(
20
- self,
21
- input_timeout: float = 60.0,
22
- print_function: Optional[Callable[..., None]] = None,
23
- on_prompt_input: Optional[Callable[[str], None]] = None,
24
- ) -> None:
25
- """
26
- Initialize the IOStream.
27
-
28
- Parameters
29
- ----------
30
- input_timeout : float, optional
31
- The input timeout in seconds, by default 60.0.
32
- print_function : Optional[Callable[..., None]], optional
33
- The function to handle print operations, by default None.
34
- on_prompt_input : Optional[Callable[[str], None]], optional
35
- The function to call for processing input prompts, by default None.
36
-
37
- Notes
38
- -----
39
- - on_prompt_input: It does not return a string (like 'input' does).
40
- Instead, it is meant to be used to forward the prompt somewhere else
41
- (e.g., a websocket). When we get the input, we can call
42
- `waldiez_io_stream.set_input(input_data)` with the input data.
43
- """
44
- self.input_timeout = input_timeout # Timeout for input
45
- self.print_function = print_function # Custom print handler
46
- self._on_prompt_input = on_prompt_input # Custom input prompt handler
47
- self.current_input: Optional[str] = None # Store the current input
48
- self._input_event = threading.Event() # Event to signal input readiness
49
- self.allow_input = True # Flag to allow or block input setting
50
-
51
- def print(
52
- self,
53
- *objects: Any,
54
- sep: str = " ",
55
- end: str = "\n",
56
- flush: bool = False,
57
- ) -> None:
58
- """
59
- Mock the `print` function.
60
-
61
- Parameters
62
- ----------
63
- objects : Any
64
- The objects to print.
65
- sep : str, optional
66
- The separator, by default " ".
67
- end : str, optional
68
- The ending string, by default a new line.
69
- flush : bool, optional
70
- Whether to flush the output, by default False.
71
- """
72
- print_function: Callable[..., None] = self.print_function or print
73
- print_function(*objects, sep=sep, end=end, flush=flush)
74
-
75
- def input(self, prompt: str = "", *, password: bool = False) -> str:
76
- """
77
- Mock the `input` function with optional timeout handling.
78
-
79
- Parameters
80
- ----------
81
- prompt : str, optional
82
- The prompt to show, by default "".
83
- password : bool, optional
84
- Whether to hide the input as password (not used), by default False.
85
-
86
- Returns
87
- -------
88
- str
89
- The user's input or '\n' if timeout occurs.
90
- """
91
- _prompt = prompt or "Your input:"
92
- if _prompt in (">", "> "): # pragma: no cover
93
- _prompt = "Your input:"
94
- if prompt:
95
- if self._on_prompt_input:
96
- self._on_prompt_input(_prompt)
97
- self.print(_prompt, end="")
98
-
99
- # Only reset if no input is currently set
100
- # e.g. handle the case when we call set_input(..) before input(..)
101
- already_set = self._input_event.is_set()
102
- if not self._input_event.is_set():
103
- self.current_input = None # Reset previous input
104
- self._input_event.clear() # Clear the event before waiting
105
-
106
- # Wait for input or timeout
107
- if not self._input_event.wait(self.input_timeout):
108
- # Timeout occurred, return what we have so far
109
- to_return = (
110
- self.current_input if self.current_input is not None else "\n"
111
- )
112
- self.current_input = None
113
- self._input_event.clear()
114
- # if we had already set the input, return it
115
- return to_return if already_set else "\n"
116
-
117
- # Input is ready, return it
118
- to_return = (
119
- self.current_input if self.current_input is not None else "\n"
120
- )
121
- self.current_input = None
122
- self._input_event.clear() # Clear the event after waiting
123
- return to_return
124
-
125
- def set_input(self, value: str) -> None:
126
- """
127
- Set the input value and signal that input is ready.
128
-
129
- Parameters
130
- ----------
131
- value : str
132
- The value to set as input.
133
- """
134
- if self.allow_input: # Respect the allow_input flag
135
- self.current_input = value
136
- self._input_event.set() # Signal that input is ready
137
-
138
- def reset_state(self) -> None:
139
- """Reset the IOStream state for testing."""
140
- self.current_input = None
141
- self._input_event.clear()
142
- self.allow_input = True # Re-enable input setting