Relative imports

modules/foo/bar/relative.krk

@@ -0,0 +1,6 @@
+from ..extra import exportablething
+
+let myotherexport = exportablething + ' there'
+
+if __name__ == '__main__':
+    print(__package__,__name__)

modules/foo/relative.krk

@@ -0,0 +1,3 @@
+from .extra import exportablething
+
+let myexport = exportablething

src/compiler.c

@@ -40,6 +40,7 @@
 #include <kuroko/object.h>
 #include <kuroko/debug.h>
 #include <kuroko/vm.h>
+#include <kuroko/util.h>
 
 /**
  * @brief Token parser state.
@@ -2177,34 +2178,43 @@ static void raiseStatement(void) {
 	emitByte(OP_RAISE);
 }
 
-static size_t importModule(KrkToken * startOfName) {
-	consume(TOKEN_IDENTIFIER, "Expected module name after 'import'.");
-	*startOfName = parser.previous;
-	while (match(TOKEN_DOT)) {
-		if (startOfName->start + startOfName->literalWidth != parser.previous.start) {
-			error("Unexpected whitespace after module path element");
-			return 0;
-		}
-		startOfName->literalWidth += parser.previous.literalWidth;
-		startOfName->length += parser.previous.length;
-		consume(TOKEN_IDENTIFIER, "Expected module path element after '.'");
-		if (startOfName->start + startOfName->literalWidth != parser.previous.start) {
-			error("Unexpected whitespace after '.'");
-			return 0;
-		}
-		startOfName->literalWidth += parser.previous.literalWidth;
-		startOfName->length += parser.previous.length;
+static size_t importModule(KrkToken * startOfName, int leadingDots) {
+	size_t ind = 0;
+	struct StringBuilder sb = {0};
+
+	for (int i = 0; i < leadingDots; ++i) {
+		pushStringBuilder(&sb, '.');
 	}
-	size_t ind = identifierConstant(startOfName);
+
+	if (!(leadingDots && check(TOKEN_IMPORT))) {
+		consume(TOKEN_IDENTIFIER, "Expected module name after 'import'.");
+		if (parser.hadError) goto _freeImportName;
+		pushStringBuilderStr(&sb, parser.previous.start, parser.previous.length);
+
+		while (match(TOKEN_DOT)) {
+			pushStringBuilderStr(&sb, parser.previous.start, parser.previous.length);
+			consume(TOKEN_IDENTIFIER, "Expected module path element after '.'");
+			if (parser.hadError) goto _freeImportName;
+			pushStringBuilderStr(&sb, parser.previous.start, parser.previous.length);
+		}
+	}
+
+	startOfName->start  = sb.bytes;
+	startOfName->length = sb.length;
+
+	ind = identifierConstant(startOfName);
 	EMIT_OPERAND_OP(OP_IMPORT, ind);
+
+_freeImportName:
+	discardStringBuilder(&sb);
 	return ind;
 }
 
 static void importStatement(void) {
 	do {
 		KrkToken firstName = parser.current;
-		KrkToken startOfName;
-		size_t ind = importModule(&startOfName);
+		KrkToken startOfName = {0};
+		size_t ind = importModule(&startOfName, 0);
 		if (match(TOKEN_AS)) {
 			consume(TOKEN_IDENTIFIER, "Expected identifier after 'as'.");
 			ind = identifierConstant(&parser.previous);
@@ -2228,8 +2238,14 @@ static void importStatement(void) {
 
 static void fromImportStatement(void) {
 	int expectCloseParen = 0;
-	KrkToken startOfName;
-	importModule(&startOfName);
+	KrkToken startOfName = {0};
+	int leadingDots = 0;
+
+	while (match(TOKEN_DOT)) {
+		leadingDots++;
+	}
+
+	importModule(&startOfName, leadingDots);
 	consume(TOKEN_IMPORT, "Expected 'import' after module name");
 	if (match(TOKEN_LEFT_PAREN)) {
 		expectCloseParen = 1;

src/kuroko/util.h

@@ -112,7 +112,7 @@ static inline void pushStringBuilder(struct StringBuilder * sb, char c) {
  * @param str C string to add.
  * @param len Length of the C string.
  */
-static inline void pushStringBuilderStr(struct StringBuilder * sb, char *str, size_t len) {
+static inline void pushStringBuilderStr(struct StringBuilder * sb, const char *str, size_t len) {
 	if (sb->capacity < sb->length + len) {
 		size_t prevcap = sb->capacity;
 		while (sb->capacity < sb->length + len) {

src/kuroko/vm.h

@@ -732,9 +732,10 @@ extern KrkValue krk_dirObject(int argc, KrkValue argv[], int hasKw);
  * @param path      Dotted path of the module, used for file lookup.
  * @param moduleOut Receives a value with the module object.
  * @param runAs     Name to attach to @c \__name__ for this module, different from @p path
+ * @param parent    Parent module object, if loaded from a package.
  * @return 1 if the module was loaded, 0 if an @ref ImportError occurred.
  */
-extern int krk_loadModule(KrkString * path, KrkValue * moduleOut, KrkString * runAs);
+extern int krk_loadModule(KrkString * path, KrkValue * moduleOut, KrkString * runAs, KrkValue parent);
 
 /**
  * @brief Load a module by a dotted name.

src/vm.c

@@ -1706,7 +1706,7 @@ static int handleException() {
  * a later search path has a krk source and an earlier search path has a shared
  * object module, the later search path will still win.
  */
-int krk_loadModule(KrkString * path, KrkValue * moduleOut, KrkString * runAs) {
+int krk_loadModule(KrkString * path, KrkValue * moduleOut, KrkString * runAs, KrkValue parent) {
 	KrkValue modulePaths;
 
 	/* See if the module is already loaded */
@@ -1771,7 +1771,20 @@ int krk_loadModule(KrkString * path, KrkValue * moduleOut, KrkString * runAs) {
 		 * returns to the current call frame; modules should return objects. */
 		KrkInstance * enclosing = krk_currentThread.module;
 		krk_startModule(runAs->chars);
-		if (isPackage) krk_attachNamedValue(&krk_currentThread.module->fields,"__ispackage__",BOOLEAN_VAL(1));
+		if (isPackage) {
+			krk_attachNamedValue(&krk_currentThread.module->fields,"__ispackage__",BOOLEAN_VAL(1));
+			/* For a module that is a package, __package__ is its own name */
+			krk_attachNamedValue(&krk_currentThread.module->fields,"__package__",OBJECT_VAL(runAs));
+		} else {
+			KrkValue parentName;
+			if (IS_INSTANCE(parent) && krk_tableGet_fast(&AS_INSTANCE(parent)->fields, S("__name__"), &parentName) && IS_STRING(parentName)) {
+				krk_attachNamedValue(&krk_currentThread.module->fields, "__package__", parentName);
+			} else {
+				/* If there is no parent, or the parent doesn't have a string __name__ attribute,
+				 * set the __package__ to None, so it at least exists. */
+				krk_attachNamedValue(&krk_currentThread.module->fields, "__package__", NONE_VAL());
+			}
+		}
 		krk_callfile(fileName,fileName);
 		*moduleOut = OBJECT_VAL(krk_currentThread.module);
 		krk_currentThread.module = enclosing;
@@ -1891,7 +1904,105 @@ int krk_importModule(KrkString * name, KrkString * runAs) {
 
 	if (isClear) {
 		KrkValue base;
-		return krk_loadModule(name,&base,runAs);
+		return krk_loadModule(name,&base,runAs,NONE_VAL());
+	}
+
+	if (name->chars[0] == '.') {
+		/**
+		 * For relative imports, we canonicalize the import name based on the current package,
+		 * and then trying importModule again with the fully qualified name.
+		 */
+
+		KrkValue packageName;
+		if (!krk_tableGet_fast(&krk_currentThread.module->fields, S("__package__"), &packageName) || !IS_STRING(packageName)) {
+			/* We must have __package__ set to a string for this to make any sense. */
+			krk_runtimeError(vm.exceptions->importError, "attempted relative import without a package context");
+			return 0;
+		}
+
+		if (name->length == 1) {
+			/* from . import ... */
+			return krk_importModule(AS_STRING(packageName), AS_STRING(packageName));
+		}
+
+		if (name->chars[1] != '.') {
+			/* from .something import ... */
+			krk_push(packageName);
+			krk_push(OBJECT_VAL(name));
+			krk_addObjects();
+
+			if (krk_importModule(AS_STRING(krk_peek(0)), AS_STRING(krk_peek(0)))) {
+				krk_swap(1); /* Imported module */
+				krk_pop(); /* Name */
+				return 1;
+			}
+
+			return 0;
+		}
+
+		/**
+		 * from .. import ...
+		 *   or
+		 * from ..something import ...
+		 *
+		 * If there n dots, there are n-1 components to pop from the end of
+		 * the package name, as '..' is "go up one" and '...' is "go up two".
+		 */
+		size_t dots = 0;
+		while (name->chars[dots+1] == '.') dots++;
+
+		/* We'll split the package name is str.split(__package__,'.') */
+		krk_push(packageName);
+		krk_push(OBJECT_VAL(S(".")));
+		KrkValue components = krk_string_split(2,(KrkValue[]){krk_peek(1),krk_peek(0)}, 0);
+		if (!IS_list(components)) {
+			krk_runtimeError(vm.exceptions->importError, "internal error while calculating package path");
+			return 0;
+		}
+		krk_push(components);
+		krk_swap(2);
+		krk_pop();
+		krk_pop();
+
+		/* If there are not enough components to "go up" through, that's an error. */
+		if (AS_LIST(components)->count <= dots) {
+			krk_runtimeError(vm.exceptions->importError, "attempted relative import beyond top-level package");
+			return 0;
+		}
+
+		size_t count = AS_LIST(components)->count - dots;
+		struct StringBuilder sb = {0};
+
+		/* Now rebuild the dotted form from the remaining components... */
+		for (size_t i = 0; i < count; i++) {
+			KrkValue node = AS_LIST(components)->values[i];
+			if (!IS_STRING(node)) {
+				discardStringBuilder(&sb);
+				krk_runtimeError(vm.exceptions->importError, "internal error while calculating package path");
+				return 0;
+			}
+			pushStringBuilderStr(&sb, AS_CSTRING(node), AS_STRING(node)->length);
+			if (i + 1 != count) {
+				pushStringBuilder(&sb, '.');
+			}
+		}
+
+		krk_pop(); /* components */
+
+		if (name->chars[dots+1]) {
+			/* from ..something import ... - append '.something' */
+			pushStringBuilderStr(&sb, &name->chars[dots], name->length - dots);
+		}
+
+		krk_push(OBJECT_VAL(finishStringBuilder(&sb)));
+
+		/* Now to try to import the fully qualified module path */
+		if (krk_importModule(AS_STRING(krk_peek(0)), AS_STRING(krk_peek(0)))) {
+			krk_swap(1); /* Imported module */
+			krk_pop(); /* Name */
+			return 1;
+		}
+		return 0;
 	}
 
 	/**
@@ -1938,7 +2049,7 @@ int krk_importModule(KrkString * name, KrkString * runAs) {
 			krk_pop(); /* dot */
 			krk_pop(); /* remainder */
 			KrkValue current;
-			if (!krk_loadModule(AS_STRING(krk_currentThread.stack[argBase+1]), &current, runAs)) return 0;
+			if (!krk_loadModule(AS_STRING(krk_currentThread.stack[argBase+1]), &current, runAs, krk_currentThread.stack[argBase-1])) return 0;
 			krk_pop(); /* dot-sepaerated */
 			krk_pop(); /* slash-separated */
 			krk_push(current);
@@ -1949,7 +2060,7 @@ int krk_importModule(KrkString * name, KrkString * runAs) {
 			return 1;
 		} else {
 			KrkValue current;
-			if (!krk_loadModule(AS_STRING(krk_currentThread.stack[argBase+1]), &current, AS_STRING(krk_currentThread.stack[argBase+2]))) return 0;
+			if (!krk_loadModule(AS_STRING(krk_currentThread.stack[argBase+1]), &current, AS_STRING(krk_currentThread.stack[argBase+2]),NONE_VAL())) return 0;
 			krk_push(current);
 			if (!IS_NONE(krk_currentThread.stack[argBase-1])) {
 				krk_tableSet(&AS_INSTANCE(krk_currentThread.stack[argBase-1])->fields, krk_currentThread.stack[argBase+0], krk_peek(0));

test/testRelativeImport.krk

@@ -0,0 +1,34 @@
+# Should work
+import foo.relative
+print(foo.relative.myexport)
+
+# Should also work
+import foo.bar.relative
+print(foo.bar.relative.myotherexport)
+
+# Should not work:
+try:
+    from . import this_is_not_a_package
+except Exception as e:
+    print("failed with", type(e).__name__)
+
+# Nor should this:
+try:
+    from .tools import this_is_not_a_package
+except Exception as e:
+    print("failed with", type(e).__name__)
+
+# But if we set __package__, we can pretend we're part of a package.
+let __package__ = 'foo'
+
+from .extra import exportablething as a
+print(a)
+
+# And if we do this...
+let __package__ = 'foo.bar.bax.qux'
+
+from ....extra import exportablething as b
+print(b)
+
+from ...relative import myotherexport as c
+print(c)

test/testRelativeImport.krk.expect

@@ -0,0 +1,9 @@
+Imported foo.__init__ as foo
+hi
+Imported bar.__init__ as foo.bar
+hi there
+failed with ImportError
+failed with ImportError
+hi
+hi
+hi there