浏览代码

Add subcommands for consumption types to cli

Maarten van den Berg 5 年之前
父节点
当前提交
eb636e6ebe
共有 2 个文件被更改,包括 154 次插入57 次删除
  1. 74 1
      piket_client/cli.py
  2. 80 56
      piket_client/model.py

+ 74 - 1
piket_client/cli.py

@@ -1,5 +1,11 @@
1 1
 import click
2
-from piket_client.model import ServerStatus, NetworkError, Consumption, Settlement
2
+from piket_client.model import (
3
+    ServerStatus,
4
+    NetworkError,
5
+    Consumption,
6
+    Settlement,
7
+    ConsumptionType,
8
+)
3 9
 from prettytable import PrettyTable
4 10
 
5 11
 
@@ -99,6 +105,73 @@ def output_settlement_info(s: Settlement) -> None:
99 105
     print(table)
100 106
 
101 107
 
108
+@cli.group()
109
+def consumption_types():
110
+    pass
111
+
112
+
113
+@consumption_types.command("list")
114
+def list_consumption_types() -> None:
115
+    active = ConsumptionType.get_all(active=True)
116
+    inactive = ConsumptionType.get_all(active=False)
117
+
118
+    if isinstance(active, NetworkError) or isinstance(inactive, NetworkError):
119
+        print_error("Could not get consumption types!")
120
+        return
121
+
122
+    table = PrettyTable()
123
+    table.field_names = ["ID", "Name", "Active"]
124
+    table.sortby = "ID"
125
+
126
+    for ct in active + inactive:
127
+        table.add_row([ct.consumption_type_id, ct.name, ct.active])
128
+
129
+    print(table)
130
+
131
+
132
+@consumption_types.command("create")
133
+@click.argument("name")
134
+def create_consumption_type(name: str) -> None:
135
+    ct = ConsumptionType(name=name).create()
136
+
137
+    if not isinstance(ct, NetworkError):
138
+        print_ok(f'Created consumption type "{name}" with ID {ct.consumption_type_id}.')
139
+
140
+
141
+@consumption_types.command("activate")
142
+@click.argument("consumption_type_id", type=click.INT)
143
+def activate_consumption_type(consumption_type_id: int) -> None:
144
+    ct = ConsumptionType.get(consumption_type_id)
145
+
146
+    if isinstance(ct, NetworkError):
147
+        print_error(f"Could not get ConsumptionType: {ct.value}")
148
+        return
149
+
150
+    result = ct.set_active(True)
151
+
152
+    if not isinstance(result, NetworkError):
153
+        print_ok(
154
+            f"Consumption type {ct.consumption_type_id} ({ct.name}) is now active."
155
+        )
156
+
157
+
158
+@consumption_types.command("deactivate")
159
+@click.argument("consumption_type_id", type=click.INT)
160
+def deactivate_consumption_type(consumption_type_id: int) -> None:
161
+    ct = ConsumptionType.get(consumption_type_id)
162
+
163
+    if isinstance(ct, NetworkError):
164
+        print_error(f"Could not get ConsumptionType: {ct.value}")
165
+        return
166
+
167
+    result = ct.set_active(False)
168
+
169
+    if not isinstance(result, NetworkError):
170
+        print_ok(
171
+            f"Consumption type {ct.consumption_type_id} ({ct.name}) is now inactive."
172
+        )
173
+
174
+
102 175
 def print_ok(msg: str) -> None:
103 176
     click.echo(click.style(msg, fg="green"))
104 177
 

+ 80 - 56
piket_client/model.py

@@ -327,85 +327,84 @@ class ConsumptionType(NamedTuple):
327 327
     name: str
328 328
     consumption_type_id: Optional[int] = None
329 329
     icon: Optional[str] = None
330
+    active: bool = True
330 331
 
331
-    def create(self) -> Optional[ConsumptionType]:
332
+    def create(self) -> Union[ConsumptionType, NetworkError]:
332 333
         """ Create a new ConsumptionType from the current attributes. As tuples
333 334
         are immutable, a new ConsumptionType with the correct id is returned.
334 335
         """
335
-        req = requests.post(
336
-            urljoin(SERVER_URL, "consumption_types"),
337
-            json={"consumption_type": {"name": self.name, "icon": self.icon}},
338
-        )
339
-
340 336
         try:
341
-            data = req.json()
342
-        except ValueError:
343
-            LOG.error(
344
-                "Did not get JSON on adding ConsumptionType (%s): %s",
345
-                req.status_code,
346
-                req.content,
337
+            req = requests.post(
338
+                urljoin(SERVER_URL, "consumption_types"),
339
+                json={"consumption_type": {"name": self.name, "icon": self.icon}},
347 340
             )
348
-            return None
349 341
 
350
-        if "error" in data or req.status_code != 201:
351
-            LOG.error(
352
-                "Could not create ConsumptionType (%s): %s", req.status_code, data
353
-            )
354
-            return None
342
+            req.raise_for_status()
343
+            data = req.json()
344
+            return ConsumptionType.from_dict(data["consumption_type"])
345
+
346
+        except requests.ConnectionError as e:
347
+            LOG.exception(e)
348
+            return NetworkError.ConnectionFailure
349
+
350
+        except requests.HTTPError as e:
351
+            LOG.exception(e)
352
+            return NetworkError.HttpFailure
353
+
354
+        except ValueError as e:
355
+            LOG.exception(e)
356
+            return NetworkError.InvalidData
355 357
 
356
-        return ConsumptionType.from_dict(data["consumption_type"])
357 358
 
358 359
     @classmethod
359
-    def get(cls, consumption_type_id: int) -> Optional[ConsumptionType]:
360
+    def get(cls, consumption_type_id: int) -> Union[ConsumptionType, NetworkError]:
360 361
         """ Retrieve a ConsumptionType by id. """
361
-        req = requests.get(
362
-            urljoin(SERVER_URL, f"/consumption_types/{consumption_type_id}")
363
-        )
364
-
365 362
         try:
363
+            req = requests.get(
364
+                urljoin(SERVER_URL, f"/consumption_types/{consumption_type_id}")
365
+            )
366
+            req.raise_for_status()
366 367
             data = req.json()
367 368
 
368
-            if "error" in data:
369
-                LOG.warning(
370
-                    "Could not get consumption type %s (%s): %s",
371
-                    consumption_type_id,
372
-                    req.status_code,
373
-                    data,
374
-                )
375
-                return None
369
+        except requests.ConnectionError as e:
370
+            LOG.exception(e)
371
+            return NetworkError.ConnectionFailure
376 372
 
377
-            return cls.from_dict(data["consumption_type"])
373
+        except requests.HTTPError as e:
374
+            LOG.exception(e)
375
+            return NetworkError.HttpFailure
378 376
 
379
-        except ValueError:
380
-            LOG.error(
381
-                "Did not get JSON from server on getting consumption type (%s): %s",
382
-                req.status_code,
383
-                req.content,
384
-            )
385
-            return None
377
+        except ValueError as e:
378
+            LOG.exception(e)
379
+            return NetworkError.InvalidData
386 380
 
387
-    @classmethod
388
-    def get_all(cls) -> Optional[List[ConsumptionType]]:
389
-        """ Get all active ConsumptionTypes. """
390
-        req = requests.get(urljoin(SERVER_URL, "/consumption_types"))
381
+        return cls.from_dict(data["consumption_type"])
391 382
 
383
+    @classmethod
384
+    def get_all(cls, active: bool = True) -> Union[List[ConsumptionType], NetworkError]:
385
+        """ Get the list of ConsumptionTypes. """
392 386
         try:
387
+            req = requests.get(
388
+                urljoin(SERVER_URL, "/consumption_types"),
389
+                params={"active": int(active)},
390
+            )
391
+            req.raise_for_status()
392
+
393 393
             data = req.json()
394 394
 
395
-            if "error" in data:
396
-                LOG.warning(
397
-                    "Could not get consumption types (%s): %s", req.status_code, data
398
-                )
395
+        except requests.ConnectionError as e:
396
+            LOG.exception(e)
397
+            return NetworkError.ConnectionFailure
399 398
 
400
-            return [cls.from_dict(item) for item in data["consumption_types"]]
399
+        except requests.HTTPError as e:
400
+            LOG.exception(e)
401
+            return NetworkError.HttpFailure
401 402
 
402
-        except ValueError:
403
-            LOG.error(
404
-                "Did not get JSON from server on getting ConsumptionTypes (%s): %s",
405
-                req.status_code,
406
-                req.content,
407
-            )
408
-            return None
403
+        except ValueError as e:
404
+            LOG.exception(e)
405
+            return NetworkError.InvalidData
406
+
407
+        return [cls.from_dict(x) for x in data["consumption_types"]]
409 408
 
410 409
     @classmethod
411 410
     def from_dict(cls, data: dict) -> "ConsumptionType":
@@ -414,8 +413,33 @@ class ConsumptionType(NamedTuple):
414 413
             name=data["name"],
415 414
             consumption_type_id=data["consumption_type_id"],
416 415
             icon=data.get("icon"),
416
+            active=data["active"],
417 417
         )
418 418
 
419
+    def set_active(self, active: bool) -> Union[ConsumptionType, NetworkError]:
420
+        """Update the 'active' attribute."""
421
+        try:
422
+            req = requests.patch(
423
+                urljoin(SERVER_URL, f"/consumption_types/{self.consumption_type_id}"),
424
+                json={"consumption_type": {"active": active}},
425
+            )
426
+            req.raise_for_status()
427
+            data = req.json()
428
+
429
+        except requests.ConnectionError as e:
430
+            LOG.exception(e)
431
+            return NetworkError.ConnectionFailure
432
+
433
+        except requests.HTTPError as e:
434
+            LOG.exception(e)
435
+            return NetworkError.HttpFailure
436
+
437
+        except ValueError as e:
438
+            LOG.exception(e)
439
+            return NetworkError.InvalidData
440
+
441
+        return self.from_dict(data["consumption_type"])
442
+
419 443
 
420 444
 class Consumption(NamedTuple):
421 445
     """ Represents a stored Consumption. """