Add methods to add/remove products from categories

magento/models/__init__.py

@@ -1,5 +1,5 @@
 from .model import Model, APIResponse
-from .category import Category
 from .product import Product, MediaEntry, ProductAttribute
+from .category import Category
 from .order import Order, OrderItem
 from .invoice import Invoice, InvoiceItem

magento/models/category.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 import copy
-from . import Model
-from typing import TYPE_CHECKING, List, Optional, Set, Dict, Union
+from . import Model, Product
 from functools import cached_property
+from magento.exceptions import MagentoError
+from typing import TYPE_CHECKING, List, Optional, Set, Dict, Union
+
 
 if TYPE_CHECKING:
     from magento import Client
-    from . import Product, Order, OrderItem, Invoice
+    from . import Order, OrderItem, Invoice
 
 
 class Category(Model):
@@ -144,3 +146,54 @@ def get_invoices(self, search_subcategories: bool = False) -> Optional[Invoice |
         :param search_subcategories: if ``True``, also searches for invoices from :attr:`~.all_subcategories`
         """
         return self.client.invoices.by_category(self, search_subcategories)
+
+    def add_product(self, product: Union[str, Product], position: Optional[int] = None) -> bool:
+        """Adds a product to the category.
+
+        .. note:: This method can also be used to update the position of a product
+           that's already in the category.
+
+        :param product: the product sku or its corresponding :class:`~.Product` object
+        :param position: the product position value to use
+        :return: success status
+        """
+        url = self.data_endpoint() + "/products"
+        sku = product.encoded_sku if isinstance(product, Product) else self.encode(product)
+        payload = {
+            "productLink": {
+                "sku": sku,
+                "category_id": self.uid
+            }
+        }
+        if isinstance(position, int):
+            payload['productLink'].update({"position": position})
+
+        response = self.client.put(url, payload)
+
+        if response.ok and response.json() is True:
+            self.logger.info(f"Added {product} to {self}")
+            return True
+        else:
+            self.logger.error(
+                f"Failed to add {product} to {self}.\nMessage: {MagentoError.parse(response)}"
+            )
+            return False
+
+    def remove_product(self, product: Union[str, Product]) -> bool:
+        """Removes a product from the category.
+
+        :param product: the product sku or its corresponding :class:`~.Product` object
+        :return: success status
+        """
+        sku = product.encoded_sku if isinstance(product, Product) else self.encode(product)
+        url = f"{self.data_endpoint()}/products/{sku}"
+        response = self.client.delete(url)
+
+        if response.ok and response.json() is True:
+            self.logger.info(f'Removed {product} from {self}')
+            return True
+        else:
+            self.logger.error(
+                f'Failed to remove {product} from {self}. Message: {MagentoError.parse(response)}'
+            )
+            return False

magento/models/product.py

@@ -6,7 +6,6 @@
 
 if TYPE_CHECKING:
     from magento import Client
-    from magento.search import SearchQuery
     from . import Category, Order, OrderItem, Invoice
 
 
@@ -129,6 +128,40 @@ def update_metadata(self, metadata: dict, scope: Optional[str] = None) -> bool:
         attributes = {k: v for k, v in metadata.items() if k in ('meta_title', 'meta_keyword', 'meta_description')}
         return self.update_custom_attributes(attributes, scope)
 
+    def add_categories(self, category_ids: Union[int, str, List[int | str]]) -> bool:
+        """Adds the product to an individual or multiple categories
+
+        :param category_ids: an individual or list of category IDs to add the product to
+        """
+        if not isinstance(category_ids, list):
+            if not isinstance(category_ids, (str, int)):
+                raise TypeError(
+                    "`category_ids` must be an individual or list of integers/strings"
+                )
+            category_ids = [category_ids]
+
+        current_ids = self.custom_attributes.get('category_ids', [])
+        new_ids = [id for id in map(str, category_ids) if id not in current_ids]
+
+        return self.update_custom_attributes({"category_ids": current_ids + new_ids})
+
+    def remove_categories(self, category_ids: Union[int, str, List[int | str]]) -> bool:
+        """Removes the product from an individual or multiple categories
+
+        :param category_ids: an individual or list of category IDs to remove the product from
+        """
+        if not isinstance(category_ids, list):
+            if not isinstance(category_ids, (str, int)):
+                raise TypeError(
+                    "`category_ids` must be an individual or list of integers/strings"
+                )
+            category_ids = [category_ids]
+
+        current_ids = self.custom_attributes.get('category_ids', [])
+        new_ids = [id for id in current_ids if id not in map(str, category_ids)]
+
+        return self.update_custom_attributes({'category_ids': new_ids})
+
     def update_attributes(self, attribute_data: dict, scope: Optional[str] = None) -> bool:
         """Update top level product attributes with scoping taken into account