From 627bda3529506eaf69df0f2308b6dfdac3c21d14 Mon Sep 17 00:00:00 2001
From: Gabriel Augendre <gabriel@augendre.info>
Date: Wed, 2 Nov 2022 22:07:33 +0100
Subject: [PATCH] Add import weapons

---
 .../commands/import_harmful_states.py         |  4 +-
 .../management/commands/import_weapons.py     | 48 +++++++++++++++++++
 src/character/migrations/0032_weapon_url.py   | 18 +++++++
 src/character/migrations/max_migration.txt    |  2 +-
 src/character/models/equipment.py             |  4 +-
 tasks.py                                      |  1 +
 6 files changed, 72 insertions(+), 5 deletions(-)
 create mode 100644 src/character/management/commands/import_weapons.py
 create mode 100644 src/character/migrations/0032_weapon_url.py

diff --git a/src/character/management/commands/import_harmful_states.py b/src/character/management/commands/import_harmful_states.py
index bdc093d..a248a0f 100644
--- a/src/character/management/commands/import_harmful_states.py
+++ b/src/character/management/commands/import_harmful_states.py
@@ -14,12 +14,12 @@ class Command(BaseCommand):
         states = self.selenium.find_elements(By.CSS_SELECTOR, "tbody tr")
         for state in states:
             try:
-                self.import_race(url, state)
+                self.import_row(url, state)
             except Exception as e:
                 print(f"{type(e)}: {e}")
         self.stdout.write(f"Finished processing {len(states)} states.")
 
-    def import_race(self, url: str, state_row: WebElement) -> None:
+    def import_row(self, url: str, state_row: WebElement) -> None:
         name = state_row.find_element(By.CLASS_NAME, "views-field-name").text.strip()
         description = state_row.find_element(
             By.CLASS_NAME, "views-field-description__value"
diff --git a/src/character/management/commands/import_weapons.py b/src/character/management/commands/import_weapons.py
new file mode 100644
index 0000000..41c4bf0
--- /dev/null
+++ b/src/character/management/commands/import_weapons.py
@@ -0,0 +1,48 @@
+from django.core.management import BaseCommand
+from selenium import webdriver
+from selenium.webdriver.common.by import By
+from selenium.webdriver.remote.webelement import WebElement
+
+from character.models import Weapon
+
+
+class Command(BaseCommand):
+    def handle(self, *args, **options) -> None:
+        url = "https://www.co-drs.org/fr/ressources/equipements/armes"
+        self.setup_selenium()
+        self.selenium.get(url)
+        states = self.selenium.find_elements(By.CSS_SELECTOR, "tbody tr")
+        for state in states:
+            try:
+                self.import_row(url, state)
+            except Exception as e:
+                print(f"{type(e)}: {e}")
+        self.stdout.write(f"Finished processing {len(states)} weapons.")
+
+    def import_row(self, url: str, state_row: WebElement) -> None:
+        name = state_row.find_element(By.CLASS_NAME, "views-field-name").text.strip()
+        category = (
+            state_row.find_element(By.CLASS_NAME, "views-field-type")
+            .text.strip()
+            .lower()
+        )
+        if "distance" in category:
+            category = Weapon.Category.RANGE
+        else:
+            category = Weapon.Category.MELEE
+        damage = state_row.find_element(By.CLASS_NAME, "views-field-dmg").text.strip()
+        weapon, _ = Weapon.objects.update_or_create(
+            name=name,
+            defaults={
+                "damage": damage,
+                "special": "",
+                "category": category,
+                "url": url,
+            },
+        )
+        self.stdout.write(self.style.SUCCESS(f"Created/updated weapon {weapon}"))
+
+    def setup_selenium(self) -> None:
+        options = webdriver.FirefoxOptions()
+        options.add_argument("-headless")
+        self.selenium = webdriver.Firefox(options=options)
diff --git a/src/character/migrations/0032_weapon_url.py b/src/character/migrations/0032_weapon_url.py
new file mode 100644
index 0000000..6913e7a
--- /dev/null
+++ b/src/character/migrations/0032_weapon_url.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.1.2 on 2022-11-02 21:08
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("character", "0031_alter_harmfulstate_options"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="weapon",
+            name="url",
+            field=models.URLField(blank=True, verbose_name="url"),
+        ),
+    ]
diff --git a/src/character/migrations/max_migration.txt b/src/character/migrations/max_migration.txt
index 626f3dd..cfadb44 100644
--- a/src/character/migrations/max_migration.txt
+++ b/src/character/migrations/max_migration.txt
@@ -1 +1 @@
-0031_alter_harmfulstate_options
+0032_weapon_url
diff --git a/src/character/models/equipment.py b/src/character/models/equipment.py
index a9d6bfc..26dde3f 100644
--- a/src/character/models/equipment.py
+++ b/src/character/models/equipment.py
@@ -1,10 +1,10 @@
 from django.db import models
 from django_extensions.db.models import TimeStampedModel
 
-from common.models import UniquelyNamedModel
+from common.models import DocumentedModel, UniquelyNamedModel
 
 
-class Weapon(UniquelyNamedModel, TimeStampedModel, models.Model):
+class Weapon(UniquelyNamedModel, DocumentedModel, TimeStampedModel, models.Model):
     class Category(models.TextChoices):
         MELEE = "MEL", "corps à corps"
         RANGE = "RAN", "à distance"
diff --git a/tasks.py b/tasks.py
index fce74b0..2f61405 100644
--- a/tasks.py
+++ b/tasks.py
@@ -128,6 +128,7 @@ def import_from_co_drs(ctx):
         ctx.run("./manage.py import_paths", pty=True, echo=True)
         ctx.run("./manage.py import_capabilities", pty=True, echo=True)
         ctx.run("./manage.py import_harmful_states", pty=True, echo=True)
+        ctx.run("./manage.py import_weapons", pty=True, echo=True)
 
 
 @task(pre=[import_from_co_drs, dump_initial])