2 files changed, 45 insertions, 13 deletions
diff --git a/src/lib/core/include/str.h b/src/lib/core/include/str.h
index bb96db0..be497ed 100644
--- a/src/lib/core/include/str.h
+++ b/src/lib/core/include/str.h
@@ -293,18 +293,27 @@ int isdigit_s(const char *s);
 char *tolower_s(char *s);
 
 /**
- * Return a copy of the input string with "." characters removed
+ * Reduce a version string to the major[minor] format used by Python
  *
- * ~~~{.c}
- * char *version = strdup("1.2.3");
- * char *version_short = to_short_version(str);
- * // version_short is "123"
- * free(version_short);
+ * @code{.c}
+ * #include <stdio.h>
+ * #include "str.h"
  *
- * ~~~
+ * int main(int argc, char *argv[]) {
+ *     char python_version[] = "3.13.3"
+ *     char *python_short_version = to_short_version(python_version);  // "313"
+ *     if (!python_short_version) {
+ *         perror("unable to allocate memory for shortened python version");
+ *         return 1;
+ *     }
+ *     free(python_short_version);
+ *     return 0;
+ * }
+ * @endcode
  *
- * @param s input string
- * @return pointer to new string
+ * @param s python version string
+ * @return the shortened version string
+ * @return NULL on error
  */
 char *to_short_version(const char *s);
 
diff --git a/src/lib/core/str.c b/src/lib/core/str.c
index 1d0b268..9524886 100644
--- a/src/lib/core/str.c
+++ b/src/lib/core/str.c
@@ -640,12 +640,35 @@ char *tolower_s(char *s) {
 }
 
 char *to_short_version(const char *s) {
-    char *result = strdup(s);
-    if (!result) {
-        return NULL;
+    char *result = NULL;
+    if (num_chars(s, '.') > 1) {
+        char **version_data = split((char *) s, ".", 1);
+        if (!version_data) {
+            goto to_short_version_failed;
+        }
+        if (version_data[1]) {
+            char *dot = strchr(version_data[1], '.');
+            if (dot) {
+                *dot = '\0';
+            }
+        }
+        result = join(version_data, "");
+        if (!result) {
+            guard_array_free(version_data);
+            goto to_short_version_failed;
+        }
+        guard_array_free(version_data);
+    } else {
+        result = strdup(s);
+        if (!result) {
+            goto to_short_version_failed;
+        }
+        strchrdel(result, ".");
     }
-    strchrdel(result, ".");
+
     return result;
+    to_short_version_failed:
+    return NULL;
 }
 
 void unindent(char *s) {