From 23601a8771f6e98e04bb4bfb9dd75ed399a20717 Mon Sep 17 00:00:00 2001
From: Ian Fan <ianfan0@gmail.com>
Date: Sun, 12 Aug 2018 00:38:19 +0100
Subject: [PATCH] commands: complete assign command

---
 include/sway/criteria.h |  9 +++++----
 sway/commands/assign.c  | 27 +++++++++++++++++++--------
 sway/sway.5.scd         |  9 ++++++++-
 sway/tree/view.c        | 23 +++++++++++++----------
 4 files changed, 45 insertions(+), 23 deletions(-)

diff --git a/include/sway/criteria.h b/include/sway/criteria.h
index b4ff7d497..7a1e547b4 100644
--- a/include/sway/criteria.h
+++ b/include/sway/criteria.h
@@ -7,10 +7,11 @@
 #include "tree/view.h"
 
 enum criteria_type {
-	CT_COMMAND          = 1 << 0,
-	CT_ASSIGN_OUTPUT    = 1 << 1,
-	CT_ASSIGN_WORKSPACE = 1 << 2,
-	CT_NO_FOCUS         = 1 << 3,
+	CT_COMMAND                 = 1 << 0,
+	CT_ASSIGN_OUTPUT           = 1 << 1,
+	CT_ASSIGN_WORKSPACE        = 1 << 2,
+	CT_ASSIGN_WORKSPACE_NUMBER = 1 << 3,
+	CT_NO_FOCUS                = 1 << 4,
 };
 
 struct criteria {
diff --git a/sway/commands/assign.c b/sway/commands/assign.c
index 0bc0929a2..04582e881 100644
--- a/sway/commands/assign.c
+++ b/sway/commands/assign.c
@@ -22,27 +22,38 @@ struct cmd_results *cmd_assign(int argc, char **argv) {
 		return error;
 	}
 
-	++argv;
-	int target_len = argc - 1;
+	--argc; ++argv;
 
 	if (strncmp(*argv, "→", strlen("→")) == 0) {
-		if (argc < 3) {
+		if (argc < 2) {
 			free(criteria);
 			return cmd_results_new(CMD_INVALID, "assign", "Missing workspace");
 		}
+		--argc;
 		++argv;
-		--target_len;
 	}
 
 	if (strcmp(*argv, "output") == 0) {
 		criteria->type = CT_ASSIGN_OUTPUT;
-		++argv;
-		--target_len;
+		--argc; ++argv;
 	} else {
-		criteria->type = CT_ASSIGN_WORKSPACE;
+		if (strcmp(*argv, "workspace") == 0) {
+			--argc; ++argv;
+		}
+		if (strcmp(*argv, "number") == 0) {
+			--argc; ++argv;
+			if (argv[0][0] < '0' || argv[0][0] > '9') {
+				free(criteria);
+				return cmd_results_new(CMD_INVALID, "assign",
+						"Invalid workspace number '%s'", argv[0]);
+			}
+			criteria->type = CT_ASSIGN_WORKSPACE_NUMBER;
+		} else {
+			criteria->type = CT_ASSIGN_WORKSPACE;
+		}
 	}
 
-	criteria->target = join_args(argv, target_len);
+	criteria->target = join_args(argv, argc);
 
 	list_add(config->criteria, criteria);
 	wlr_log(WLR_DEBUG, "assign: '%s' -> '%s' added", criteria->raw,
diff --git a/sway/sway.5.scd b/sway/sway.5.scd
index 8e56d5bb7..831880671 100644
--- a/sway/sway.5.scd
+++ b/sway/sway.5.scd
@@ -219,13 +219,20 @@ They are expected to be used with *bindsym* or at runtime through *swaymsg*(1).
 The following commands may be used either in the configuration file or at
 runtime.
 
-*assign* <criteria> [→] <workspace>
+*assign* <criteria> [→] [workspace] [number] <workspace>
 	Assigns views matching _criteria_ (see *CRITERIA* for details) to
 	_workspace_. The → (U+2192) is optional and cosmetic. This command is
 	equivalent to:
 
 		for\_window <criteria> move container to workspace <workspace>
 
+*assign* <criteria> [→] output left|right|up|down|<name>
+	Assigns views matching _criteria_ (see *CRITERIA* for details) to the
+	specified output. The → (U+2192) is optional and cosmetic. This command is
+	equivalent to:
+
+		for\_window <criteria> move container to output <output>
+
 *bindsym* [--release|--locked] <key combo> <command>
 	Binds _key combo_ to execute the sway command _command_ when pressed. You
 	may use XKB key names here (*xev*(1) is a good tool for discovering these).
diff --git a/sway/tree/view.c b/sway/tree/view.c
index 4495c1507..4c8e17747 100644
--- a/sway/tree/view.c
+++ b/sway/tree/view.c
@@ -450,12 +450,22 @@ static struct sway_container *select_workspace(struct sway_view *view) {
 
 	// Check if there's any `assign` criteria for the view
 	list_t *criterias = criteria_for_view(view,
-			CT_ASSIGN_WORKSPACE | CT_ASSIGN_OUTPUT);
+			CT_ASSIGN_WORKSPACE | CT_ASSIGN_WORKSPACE_NUMBER | CT_ASSIGN_OUTPUT);
 	struct sway_container *ws = NULL;
 	for (int i = 0; i < criterias->length; ++i) {
 		struct criteria *criteria = criterias->items[i];
-		if (criteria->type == CT_ASSIGN_WORKSPACE) {
-			ws = workspace_by_name(criteria->target);
+		if (criteria->type == CT_ASSIGN_OUTPUT) {
+			struct sway_container *output = output_by_name(criteria->target);
+			if (output) {
+				ws = seat_get_active_child(seat, output);
+				break;
+			}
+		} else {
+			// CT_ASSIGN_WORKSPACE(_NUMBER)
+			ws = criteria->type == CT_ASSIGN_WORKSPACE_NUMBER ?
+				workspace_by_number(criteria->target) :
+				workspace_by_name(criteria->target);
+
 			if (!ws) {
 				if (strcasecmp(criteria->target, "back_and_forth") == 0) {
 					if (prev_workspace_name) {
@@ -466,13 +476,6 @@ static struct sway_container *select_workspace(struct sway_view *view) {
 				}
 			}
 			break;
-		} else {
-			// CT_ASSIGN_OUTPUT
-			struct sway_container *output = output_by_name(criteria->target);
-			if (output) {
-				ws = seat_get_active_child(seat, output);
-				break;
-			}
 		}
 	}
 	list_free(criterias);