Coverage for jstark / sample / mealkit_orders.py: 100%

1import random 

2from functools import cached_property 

3import uuid 

4from datetime import date 

5from typing import Any, Iterable 

6from decimal import Decimal 

7from pyspark.sql import SparkSession, DataFrame 

8from pyspark.sql.types import ( 

9 IntegerType, 

10 StringType, 

11 StructField, 

12 StructType, 

13 TimestampType, 

14 DecimalType, 

15) 

16from faker import Faker 

17from faker.providers import DynamicProvider 

18 

19 

20class FakeMealkitOrders: 

21 def __init__(self, seed: int | None = None, number_of_orders: int = 1000): 

22 self.seed = seed 

23 self.number_of_orders = number_of_orders 

24 

25 @property 

26 def mealkit_orders_schema(self) -> StructType: 

27 return StructType( 

28 [ 

29 StructField("Timestamp", TimestampType(), True), 

30 StructField("Customer", StringType(), True), 

31 StructField("Product", StringType(), True), 

32 StructField("Recipe", StringType(), True), 

33 StructField("Cuisine", StringType(), True), 

34 StructField("Allergen", StringType(), True), 

35 StructField("Quantity", IntegerType(), True), 

36 StructField("Order", StringType(), True), 

37 StructField("Discount", DecimalType(10, 2), True), 

38 ] 

39 ) 

40 

41 @staticmethod 

42 def flatten_mealkit_orders(mealkit_orders: list[Any]) -> Iterable[dict[str, Any]]: 

43 return [ 

44 { 

45 "Customer": d["Customer"], 

46 "Product": d["Product"], 

47 "Cuisine": d["Cuisine"], 

48 "Allergen": d["Allergen"], 

49 "Order": d["Order"], 

50 "Timestamp": d["Timestamp"], 

51 **d2, 

52 } 

53 for d in mealkit_orders 

54 for d2 in d["Recipes"] 

55 ] 

56 

57 @cached_property 

58 def df(self) -> DataFrame: 

59 products_provider = DynamicProvider( 

60 provider_name="product", 

61 elements=[ 

62 "classic-plan", 

63 "preset-box-bc", 

64 "balanced-living-t1", 

65 "dinner-box", 

66 "classic-plan-t11", 

67 "classic-plan-t12", 

68 "classic-plan-t13", 

69 "classic-plan-t14", 

70 "classic-plan-t15", 

71 ], 

72 ) 

73 recipes_provider = DynamicProvider( 

74 provider_name="recipe", 

75 elements=[ 

76 "Banging bangers and mash", 

77 "Fish and chips", 

78 "Pizza and salad", 

79 "Chicken curry", 

80 "Beef stew", 

81 "Vegetable lasagna", 

82 "Salad and bread", 

83 "Soup and bread", 

84 "Pasta and sauce", 

85 ], 

86 ) 

87 cuisines_provider = DynamicProvider( 

88 provider_name="cuisine", 

89 elements=[ 

90 "Italian", 

91 "French", 

92 "Spanish", 

93 ], 

94 ) 

95 

96 allergens_provider = DynamicProvider( 

97 provider_name="allergen", 

98 elements=[ 

99 "Gluten", 

100 "Eggs", 

101 "Milk", 

102 "Soy", 

103 ], 

104 ) 

105 

106 fake = Faker() 

107 if self.seed: 

108 Faker.seed(self.seed) 

109 

110 products_fake = Faker() 

111 products_fake.add_provider(products_provider) 

112 

113 recipes_fake = Faker() 

114 recipes_fake.add_provider(recipes_provider) 

115 

116 cuisines_fake = Faker() 

117 cuisines_fake.add_provider(cuisines_provider) 

118 

119 allergens_fake = Faker() 

120 allergens_fake.add_provider(allergens_provider) 

121 

122 mealkit_orders = [] 

123 

124 possible_quantities = [1, 2] 

125 if self.seed: 

126 random.seed(self.seed) 

127 quantities = random.choices( 

128 possible_quantities, 

129 weights=[100, 1], 

130 k=self.number_of_orders * len(recipes_provider.elements), 

131 ) 

132 for order in range(self.number_of_orders): 

133 recipes = [] 

134 # if self.seed: 

135 # random.seed(self.seed) 

136 for recipe_index in range(random.randint(2, 5)): 

137 r = recipes_fake.unique.recipe() 

138 quantity = quantities[(order * len(possible_quantities)) + recipe_index] 

139 recipes.append( 

140 { 

141 "Recipe": r, 

142 "Quantity": quantity, 

143 } 

144 ) 

145 mealkit_orders.append( 

146 { 

147 "Customer": fake.name(), 

148 "Timestamp": fake.date_time_between( 

149 start_date=date(2021, 1, 1), end_date=date(2021, 12, 31) 

150 ), 

151 "Order": str(uuid.uuid4()), 

152 "Product": products_fake.product(), 

153 "Cuisine": cuisines_fake.cuisine(), 

154 "Allergen": allergens_fake.allergen(), 

155 "Recipes": recipes, 

156 "Discount": Decimal(random.uniform(0, 5)), 

157 } 

158 ) 

159 recipes_fake.unique.clear() 

160 cuisines_fake.unique.clear() 

161 flattened_mealkit_orders = self.flatten_mealkit_orders(mealkit_orders) 

162 spark = SparkSession.builder.getOrCreate() 

163 return spark.createDataFrame( 

164 flattened_mealkit_orders, 

165 schema=self.mealkit_orders_schema, # type: ignore 

166 )