aggregate(): add colnames to output in the matrix case

Author: AllinCottrell

lib/src/genfuncs.c

@@ -7319,60 +7319,76 @@ static gretl_matrix *real_aggregate_by (const double *x,
 */
 
 static void aggr_add_colnames (gretl_matrix *m,
-                               const int *ylist,
                                const int *xlist,
-                               const DATASET *dset,
-                               int just_count)
+                               const int *ylist,
+                               const gretl_matrix *X,
+                               const gretl_matrix *Y,
+                               const DATASET *dset)
 {
     char **S = NULL;
+    char **Sl = NULL;
+    const char **Sm;
     int i, j, n = m->cols;
+    int ny = 0;
+    int nx = 0;
     int err = 0;
 
     S = strings_array_new(n);
     if (S == NULL) {
         return;
     }
 
+    ny = ylist != NULL ? ylist[0] : Y->cols;
+    nx = n - ny - 1;
     j = 0;
 
-    if (ylist != NULL && ylist[0] > 0) {
-        char **Sy = gretl_list_get_names_array(ylist, dset, &err);
-
-        if (!err) {
-            for (i=0; i<ylist[0]; i++) {
-                S[j++] = Sy[i];
-                Sy[i] = NULL;
-            }
-            free(Sy);
+    if ((Sl = gretl_list_get_names_array(ylist, dset, NULL)) != NULL) {
+        for (i=0; i<ny; i++) {
+            S[j++] = Sl[i];
+            Sl[i] = NULL;
+        }
+        free(Sl);
+    } else if ((Sm = gretl_matrix_get_colnames(Y)) != NULL) {
+        for (i=0; i<ny; i++) {
+            S[j++] = gretl_strdup(Sm[i]);
         }
     } else {
-        S[j++] = gretl_strdup("byvar");
-    }
-
-    if (!err) {
-        S[j++] = gretl_strdup("count");
+        for (i=0; i<ny; i++) {
+            S[j++] = gretl_strdup("byvar");
+        }
     }
 
-    if (!err && j < n) {
-        if (xlist != NULL && xlist[0] > 0) {
-            char **Sx = gretl_list_get_names_array(xlist, dset, &err);
+    S[j++] = gretl_strdup("count");
 
-            if (!err) {
-                for (i=0; i<xlist[0]; i++) {
-                    S[j++] = Sx[i];
-                    Sx[i] = NULL;
-                }
-                free(Sx);
+    if (nx > 0) {
+        if ((Sl = gretl_list_get_names_array(xlist, dset, NULL)) != NULL) {
+            for (i=0; i<nx; i++) {
+                S[j++] = Sl[i];
+                Sl[i] = NULL;
+            }
+            free(Sl);
+        } else if ((Sm = gretl_matrix_get_colnames(X)) != NULL) {
+            for (i=0; i<nx; i++) {
+                S[j++] = gretl_strdup(Sm[i]);
             }
         } else {
-            S[j] = gretl_strdup("f(x)");
+            for (i=0; i<nx; i++) {
+                S[j++] = gretl_strdup("f(x)");
+            }
+        }
+    }
+
+    /* basic check on validity of @S */
+    for (i=0; i<n; i++) {
+        if (S[i] == NULL) {
+            strings_array_free(S, n);
+            err = 1;
+            break;
         }
     }
 
     if (!err) {
         gretl_matrix_set_colnames(m, S);
-    } else {
-        strings_array_free(S, n);
     }
 }
 
@@ -7557,7 +7573,7 @@ gretl_matrix *aggregate_by (const double *x,
     }
 
     if (m != NULL) {
-        aggr_add_colnames(m, ylist, xlist, dset, just_count);
+        aggr_add_colnames(m, xlist, ylist, NULL, NULL, dset);
     }
 
     free(tmp);
@@ -7747,6 +7763,8 @@ gretl_matrix *matrix_aggregate (const gretl_matrix *X,
         }
     }
 
+    aggr_add_colnames(ret, NULL, NULL, X, Y, NULL);
+
     gretl_matrix_free(M);
     free(counts);
     free(cols);

lib/src/gretl_list.c

@@ -1080,7 +1080,7 @@ char *gretl_list_get_names (const int *list, const DATASET *dset,
  * gretl_list_get_names_array:
  * @list: array of integers.
  * @dset: dataset information.
- * @err: location to receive error code.
+ * @err: location to receive error code, or #NULL.
  *
  * Returns: An array of strings holding the names of the
  * members of @list, or NULL on failure.
@@ -1093,8 +1093,10 @@ char **gretl_list_get_names_array (const int *list,
     char **S = NULL;
     int i, vi, n;
 
-    if (list == NULL) {
-	*err = E_DATA;
+    if (list == NULL || dset == NULL) {
+        if (err != NULL) {
+            *err = E_DATA;
+        }
 	return NULL;
     }
 
@@ -1106,7 +1108,9 @@ char **gretl_list_get_names_array (const int *list,
 
     S = strings_array_new(n);
     if (S == NULL) {
-	*err = E_ALLOC;
+        if (err != NULL) {
+            *err = E_ALLOC;
+        }
 	return NULL;
     }
 
@@ -1118,7 +1122,9 @@ char **gretl_list_get_names_array (const int *list,
 	    S[i] = gretl_strdup(dset->varname[vi]);
 	}
 	if (S[i] == NULL) {
-	    *err = E_ALLOC;
+            if (err != NULL) {
+                *err = E_ALLOC;
+            }
 	    strings_array_free(S, n);
 	    S = NULL;
 	    break;