Skip to content

Commit

Permalink
fix: update collection extender
Browse files Browse the repository at this point in the history
  • Loading branch information
adhtruong committed Aug 4, 2024
1 parent d75b1d2 commit a8bb48d
Showing 1 changed file with 7 additions and 15 deletions.
22 changes: 7 additions & 15 deletions polyfactory/collection_extender.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
from __future__ import annotations

import random
from abc import ABC, abstractmethod
from collections import deque
from typing import Any
Expand Down Expand Up @@ -54,36 +53,29 @@ class ListLikeExtender(CollectionExtender):
__types__ = (list, deque)

@staticmethod
def _extend_type_args(type_args: tuple[Any, ...], number_of_args: int) -> tuple[Any, ...]:
if not type_args:
return type_args
return tuple(random.choice(type_args) for _ in range(number_of_args))
def _extend_type_args(type_args: tuple[Any, ...], _: int) -> tuple[Any, ...]:
return type_args


class SetExtender(CollectionExtender):
__types__ = (set, frozenset)

@staticmethod
def _extend_type_args(type_args: tuple[Any, ...], number_of_args: int) -> tuple[Any, ...]:
if not type_args:
return type_args
return tuple(random.choice(type_args) for _ in range(number_of_args))
def _extend_type_args(type_args: tuple[Any, ...], _: int) -> tuple[Any, ...]:
return type_args


class DictExtender(CollectionExtender):
__types__ = (dict,)

@staticmethod
def _extend_type_args(type_args: tuple[Any, ...], number_of_args: int) -> tuple[Any, ...]:
return type_args * number_of_args
def _extend_type_args(type_args: tuple[Any, ...], _: int) -> tuple[Any, ...]:
return type_args


class FallbackExtender(CollectionExtender):
__types__ = ()

@staticmethod
def _extend_type_args(
type_args: tuple[Any, ...],
number_of_args: int, # noqa: ARG004
) -> tuple[Any, ...]: # - investigate @guacs
def _extend_type_args(type_args: tuple[Any, ...], _: int) -> tuple[Any, ...]:
return type_args

0 comments on commit a8bb48d

Please sign in to comment.